Bump jax to 0.7.1 (#2134)

**Context:**

**Description of the Change:**

1. API change of `lower_jaxpr_to_fun`, update it in catalyst accordingly
2. API change of `jaxpr_subcomp`, update it in catalyst accordingly
3. `_jaxpr_replicas` has been moved to `jax._src.interpreters.pxla` in
jax
4. After jax 0.7.1+ `mlir.multi_broadcast_in_dim` has changed, and break
the dynamic shape, patch it to fix the dynamic shape issue

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

Co-authored-by: JerryChen97 <[email protected]>
Co-authored-by: Paul <[email protected]>
Co-authored-by: David Ittah <[email protected]>
Co-authored-by: Ali Asadi <[email protected]>

Author: 5 people

.dep-versions

@@ -1,7 +1,7 @@
 # Always update the version check in catalyst.__init__ when changing the JAX version.
 # To update JAX version alongside compatible dependency tags, run the following script:
 # python3 .github/workflows/set_dep_versions.py {JAX_version}
-jax=0.7.0
+jax=0.7.1
 stablehlo=0a4440a5c8de45c4f9649bf3eb4913bf3f97da0d
 llvm=113f01aa82d055410f22a9d03b3468fa68600589
 enzyme=v0.0.203

frontend/catalyst/__init__.py

@@ -22,7 +22,7 @@
 
 import jaxlib as _jaxlib
 
-_jaxlib_version = "0.7.0"
+_jaxlib_version = "0.7.1"
 if _jaxlib.__version__ != _jaxlib_version:
     import warnings
 

frontend/catalyst/jax_extras/lowering.py

@@ -19,9 +19,10 @@
 import textwrap
 
 import jax
-from jax._src.dispatch import jaxpr_replicas
+from jax._src import core
 from jax._src.effects import ordered_effects as jax_ordered_effects
 from jax._src.interpreters.mlir import _module_name_regex
+from jax._src.interpreters.pxla import _jaxpr_replicas as jaxpr_replicas
 from jax._src.sharding_impls import AxisEnv, ReplicaAxisContext
 from jax.extend.core import ClosedJaxpr
 from jax.interpreters.mlir import (
@@ -44,7 +45,11 @@
 
 __all__ = ("jaxpr_to_mlir", "custom_lower_jaxpr_to_module")
 
-from catalyst.jax_extras.patches import _no_clean_up_dead_vars, get_aval2
+from catalyst.jax_extras.patches import (
+    _no_clean_up_dead_vars,
+    get_aval2,
+    patched_multi_broadcast_in_dim,
+)
 
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
@@ -67,6 +72,7 @@ def jaxpr_to_mlir(jaxpr, func_name, arg_names):
     with Patcher(
         (jax._src.interpreters.partial_eval, "get_aval", get_aval2),
         (jax._src.core, "clean_up_dead_vars", _no_clean_up_dead_vars),
+        (jax._src.interpreters.mlir, "multi_broadcast_in_dim", patched_multi_broadcast_in_dim),
     ):
         nrep = jaxpr_replicas(jaxpr)
         effects = jax_ordered_effects.filter_in(jaxpr.effects)
@@ -141,13 +147,20 @@ def custom_lower_jaxpr_to_module(
         module_name = _module_name_regex.sub("_", module_name)
         ctx.module.operation.attributes["sym_name"] = ir.StringAttr.get(module_name)
 
+        const_args = core.jaxpr_const_args(jaxpr.jaxpr)
+        const_arg_avals = [core.shaped_abstractify(c) for c in const_args]
+        num_const_args = len(const_arg_avals)
+        in_avals = const_arg_avals + jaxpr.in_avals
+
         # Use main_function=False to preserve the function name (e.g., "jit_func")
         # instead of renaming it to "main"
         lower_jaxpr_to_fun(
             ctx,
             func_name,
             jaxpr,
             effects,
+            num_const_args=num_const_args,
+            in_avals=in_avals,
             main_function=False,
             replicated_args=replicated_args,
             arg_names=arg_names,

frontend/catalyst/jax_extras/patches.py

@@ -23,6 +23,7 @@
 import jax._src.interpreters.partial_eval as pe
 from jax._src import config, core, source_info_util
 from jax._src.core import JaxprEqnContext, abstractify, standard_vma_rule
+from jax._src.interpreters import mlir
 from jax._src.interpreters.partial_eval import (
     DynamicJaxprTracer,
     TracingEqn,
@@ -54,6 +55,7 @@
     "patched_make_eqn",
     "patched_dyn_shape_staging_rule",
     "patched_pjit_staging_rule",
+    "patched_multi_broadcast_in_dim",
 )
 
 
@@ -322,6 +324,34 @@ def make_eqn_internal(out_avals_list, out_tracers):
         return make_eqn_internal(out_avals, out_tracers)
 
 
+def patched_multi_broadcast_in_dim(ctx, ops, ops_avals, out_shape, out_sharding=None):
+    """Patched version that uses DShapedArray for dynamic shapes."""
+    out = []
+    for op, op_aval in zip(ops, ops_avals):
+        op_aval_shape = op_aval.shape
+
+        # Use DShapedArray if shape contains dynamic dimensions
+        if core.is_constant_shape(out_shape):
+            out_aval = core.ShapedArray(out_shape, op_aval.dtype, sharding=out_sharding)
+        else:
+            # DShapedArray doesn't support sharding parameter
+            out_aval = core.DShapedArray(
+                out_shape, op_aval.dtype, weak_type=getattr(op_aval, "weak_type", False)
+            )
+
+        if core.definitely_equal_shape(op_aval_shape, out_shape):
+            out.append(op)
+        else:
+            assert len(op_aval_shape) <= len(out_shape), (op_aval_shape, out_shape)
+            broadcast_dimensions = list(range(len(out_shape) - len(op_aval_shape), len(out_shape)))
+            b_out = mlir.broadcast_in_dim(
+                ctx, op, out_aval, broadcast_dimensions=broadcast_dimensions
+            )
+            b_out = mlir.lower_with_sharding_in_types(ctx, b_out, out_aval)
+            out.append(b_out)
+    return out
+
+
 def patched_dyn_shape_staging_rule(trace, source_info, prim, out_aval, *args, **params):
     """Patched _dyn_shape_staging_rule for dynamic shape handling."""
     eqn, out_tracer = trace.make_eqn(args, out_aval, prim, params, core.no_effects, source_info)

frontend/catalyst/jax_primitives.py

@@ -2243,6 +2243,7 @@ def emit_branches(preds, branch_jaxprs, ip):
                     [mlir.ir_constant(c) for c in true_jaxpr.consts],  # is never hit in our tests
                     *flat_args_plus_consts,
                     dim_var_values=jax_ctx.dim_var_values,
+                    const_lowering=jax_ctx.const_lowering,
                 )
 
                 YieldOp(out)
@@ -2263,6 +2264,7 @@ def emit_branches(preds, branch_jaxprs, ip):
                         [mlir.ir_constants(c) for c in otherwise_jaxpr.consts],
                         *flat_args_plus_consts,
                         dim_var_values=jax_ctx.dim_var_values,
+                        const_lowering=jax_ctx.const_lowering,
                     )
 
                     YieldOp(out)
@@ -2333,6 +2335,7 @@ def _switch_lowering(
                 [mlir.ir_constant(const) for const in branch_jaxpr.consts],
                 *flat_args_plus_consts,
                 dim_var_values=jax_ctx.dim_var_values,
+                const_lowering=jax_ctx.const_lowering,
             )
 
             YieldOp(out)
@@ -2348,6 +2351,7 @@ def _switch_lowering(
             [mlir.ir_constant(const) for const in branch_jaxpr.consts],
             *flat_args_plus_consts,
             dim_var_values=jax_ctx.dim_var_values,
+            const_lowering=jax_ctx.const_lowering,
         )
 
         YieldOp(out)
@@ -2440,6 +2444,7 @@ def _while_loop_lowering(
             [mlir.ir_constants(c) for c in cond_jaxpr.consts],
             *params,
             dim_var_values=jax_ctx.dim_var_values,
+            const_lowering=jax_ctx.const_lowering,
         )
 
         pred_extracted = TensorExtractOp(ir.IntegerType.get_signless(1), pred, []).result
@@ -2461,6 +2466,7 @@ def _while_loop_lowering(
             [mlir.ir_constants(c) for c in cond_jaxpr.consts],
             *params,
             dim_var_values=jax_ctx.dim_var_values,
+            const_lowering=jax_ctx.const_lowering,
         )
 
         YieldOp(out)
@@ -2591,6 +2597,7 @@ def _for_loop_lowering(
             [mlir.ir_constants(c) for c in body_jaxpr.consts],
             *loop_params,
             dim_var_values=jax_ctx.dim_var_values,
+            const_lowering=jax_ctx.const_lowering,
         )
         YieldOp(out)
 
@@ -2720,6 +2727,7 @@ def _adjoint_lowering(
             [mlir.ir_constants(c) for c in jaxpr.consts],
             *list(chain(consts, cargs, adjoint_block.arguments)),
             dim_var_values=jax_ctx.dim_var_values,
+            const_lowering=jax_ctx.const_lowering,
         )
 
         QYieldOp([out[-1]])
@@ -2734,7 +2742,7 @@ def _adjoint_lowering(
         def adjoint_pass_injector(_op: ir.Operation) -> ir.WalkResult:
             if _op.name == "transform.named_sequence":
                 with ir.InsertionPoint.at_block_begin(_op.regions[0].blocks[0]):
-                    adjoint_lowering_pass_op = ApplyRegisteredPassOp(
+                    ApplyRegisteredPassOp(
                         result=ir.OpaqueType.get("transform", 'op<"builtin.module">'),
                         target=_op.regions[0].blocks[0].arguments[0],  # just insert at beginning
                         pass_name="adjoint-lowering",

frontend/catalyst/jax_primitives_utils.py

@@ -178,6 +178,13 @@ def lower_callable_to_funcop(ctx, callable_, call_jaxpr, public=False):
     kwargs["effects"] = []
     kwargs["main_function"] = False
 
+    const_args = core.jaxpr_const_args(call_jaxpr.jaxpr)
+    const_arg_avals = [core.shaped_abstractify(c) for c in const_args]
+    num_const_args = len(const_arg_avals)
+
+    kwargs["in_avals"] = const_arg_avals + call_jaxpr.in_avals
+    kwargs["num_const_args"] = num_const_args
+
     func_op = mlir.lower_jaxpr_to_fun(**kwargs)
     if public:
         func_op.attributes["sym_visibility"] = ir.StringAttr.get("public")

frontend/catalyst/jax_tracer.py

@@ -713,7 +713,8 @@ def lower_jaxpr_to_mlir(jaxpr, func_name, arg_names):
     # JAX internally calls trace_to_jaxpr_dynamic2 during lowering of nested @jit primitives
     # (e.g., in jax.scipy.linalg.expm and jax.scipy.linalg.solve), which triggers two bugs:
     # 1. make_eqn signature changed to include out_tracers parameter
-    # 2. pjit_staging_rule creates JaxprEqn instead of TracingEqn (AssertionError at partial_eval.py:1790)
+    # 2. pjit_staging_rule creates JaxprEqn instead of TracingEqn
+    #   (AssertionError at partial_eval.py:1790)
     with transient_jax_config(
         {"jax_dynamic_shapes": True, "jax_use_shardy_partitioner": False}
     ), Patcher(