[Pallas TPU] Convert pattern_match_while_to_fori_loop to return (Jaxpr, str) rather than throw exceptions.

Currently, pattern_match_while_to_fori_loop attempts to convert a while_loop jaxpr into a type of fori_loop which Pallas can lower.
To do so, it validates the conditions which would block the jaxpr from being lowered successfully. Because Pallas presently only supports "fori convertable" loops, this matching code also throws Exceptions when the supported conditions are violated.

In the near future, we aim to have support for more ordinary while loops -- but we still would like to perform this match-and-convert procedure when possible.
To facilitate that, this updates the error handling in pattern_match_while_to_fori_loop to simply return errors when hit, so the calling code can determine if they should be thrown.

PiperOrigin-RevId: 623274837

Author: jax authors

jax/_src/pallas/mosaic/lowering.py

@@ -1822,9 +1822,12 @@ def _while_lowering_rule(
     body_nconsts,
     body_jaxpr,
 ):
-  jaxpr = pallas_utils.pattern_match_while_to_fori_loop(
+  jaxpr, err = pallas_utils.pattern_match_while_to_fori_loop(
       cond_jaxpr, cond_nconsts, body_jaxpr, body_nconsts
   )
+  if jaxpr is None:
+    raise NotImplementedError(err)
+
   _, body_consts, carry = split_list(args, [cond_nconsts, body_nconsts])
   (lb, ub), args = carry[:2], carry[2:]
   for_out = _lower_jaxpr_to_for_loop(

jax/_src/pallas/utils.py

@@ -14,6 +14,8 @@
 
 """Pallas utility functions."""
 
+from __future__ import annotations
+
 from jax import lax
 from jax._src import core as jax_core
 from jax._src.util import split_list
@@ -98,44 +100,42 @@ def pattern_match_while_to_fori_loop(
     cond_nconsts: int,
     body_jaxpr: jax_core.Jaxpr,
     body_nconsts: int,
-) -> tuple[jax_core.Jaxpr, bool]:
+) -> tuple[jax_core.Jaxpr | None, str | None]:
   # Try to pattern match to fori loop.
+  # Successful matches produce (jaxpr, None), while failures use the str
+  # component of the return tuple to capture information about the failure.
   if cond_nconsts:
-    raise NotImplementedError("Conditional jaxpr can't contain consts.")
+    return (None, "Conditional jaxpr can't contain consts.")
   _, cond_invars = split_list(cond_jaxpr.jaxpr.invars, [cond_nconsts])
   cond_in_avals = [v.aval for v in cond_invars]
   if len(cond_in_avals) < 2:
-    raise NotImplementedError("Conditional jaxpr have only two carry args.")
+    return (None, "Conditional jaxpr have only two carry args.")
   # Check that the first two carry values are scalar ints
   a1, a2 = cond_in_avals[:2]
   if a1.shape or a1.dtype not in (jnp.int32, jnp.int64):
-    raise NotImplementedError(
-        "First conditional jaxpr carry arg is not a scalar int."
-    )
+    return (None, "First conditional jaxpr carry arg is not a scalar int.")
   if a2.shape or a2.dtype not in (jnp.int32, jnp.int64):
-    raise NotImplementedError(
-        "Second conditional jaxpr carry arg is not a scalar int."
-    )
+    return (None, "Second conditional jaxpr carry arg is not a scalar int.")
   # Check that the only eqn in the cond checks the loop index condition
   v1, v2 = cond_invars[:2]
   outvar = cond_jaxpr.jaxpr.outvars[0]
   assert outvar.aval.dtype == jnp.bool_
   if len(cond_jaxpr.jaxpr.eqns) != 1:
-    raise NotImplementedError("Non-trivial conditional jaxprs not supported.")
+    return (None, "Non-trivial conditional jaxprs not supported.")
   eqn = cond_jaxpr.jaxpr.eqns[0]
   if eqn.primitive != lax.lt_p:
-    raise NotImplementedError("Non-trivial conditional jaxprs not supported.")
+    return (None, "Non-trivial conditional jaxprs not supported.")
   if eqn.outvars != [outvar]:
-    raise NotImplementedError("Non-trivial conditional jaxprs not supported.")
+    return (None, "Non-trivial conditional jaxprs not supported.")
   if eqn.invars != [v1, v2]:
-    raise NotImplementedError("Non-trivial conditional jaxprs not supported.")
+    return (None, "Non-trivial conditional jaxprs not supported.")
   # Check that the carry is updated in the body appropriately
   _, body_invars = split_list(body_jaxpr.jaxpr.invars, [body_nconsts])
   v1, v2 = body_invars[:2]
   vo1, vo2 = body_jaxpr.jaxpr.outvars[:2]
   # Upper bound should be constant
   if v2 is not vo2:
-    raise NotImplementedError("Loop upper bound is not constant.")
+    return (None, "Loop upper bound is not constant.")
   # Check that we increment the loop index in the body
   for i, eqn in enumerate(body_jaxpr.jaxpr.eqns):
     if eqn.primitive is lax.add_p:
@@ -146,7 +146,7 @@ def pattern_match_while_to_fori_loop(
               eqn_index = i
               break
   else:
-    raise NotImplementedError("Loop index not incremented in body.")
+    return (None, "Loop index not incremented in body.")
   jaxpr = body_jaxpr.jaxpr
   new_invars = (
       *jaxpr.invars[:body_nconsts],
@@ -159,4 +159,4 @@ def pattern_match_while_to_fori_loop(
       invars=new_invars,
       outvars=new_outvars,
   )
-  return jaxpr
+  return jaxpr, None