Merge pull request #33841 from mattjj:bjp

PiperOrigin-RevId: 842541857

Author: Google-ML-Automation

jax/_src/interpreters/mlir.py

@@ -2451,8 +2451,15 @@ def f_lowered(ctx: LoweringRuleContext, *args, **params):
     wrapped_fun = lu.wrap_init(f, params,
         debug_info=api_util.debug_info("lower_fun", fun, args, {}))
 
-    jaxpr, _, consts_for_constvars = pe.trace_to_jaxpr_dynamic(wrapped_fun,
-                                                               ctx.avals_in)
+    jaxpr, _, consts_for_constvars = pe.trace_to_jaxpr_dynamic(
+        wrapped_fun, ctx.avals_in)
+
+    if any(isinstance(e, core.InternalMutableArrayEffect) for e in jaxpr.effects):
+      from jax._src.interpreters import pxla  # type: ignore
+      closed_jaxpr = core.ClosedJaxpr(jaxpr, consts_for_constvars)
+      closed_jaxpr = pxla._discharge_internal_refs(closed_jaxpr)
+      jaxpr, consts_for_constvars = closed_jaxpr.jaxpr, closed_jaxpr.consts
+
     # TODO(frostig,mattjj): check ctx.avals_out against jaxpr avals out?
 
     if ctx.platforms is not None:

tests/mutable_array_test.py

@@ -27,6 +27,7 @@
 from jax._src import test_util as jtu
 from jax._src.api import vjp3
 from jax._src.util import safe_map, safe_zip
+from jax._src.interpreters import mlir
 from jax.sharding import NamedSharding, PartitionSpec as P, AxisType
 import jax.numpy as jnp
 
@@ -1031,6 +1032,28 @@ def test_none_index(self):
     y = ref[None]
     self.assertEqual(y.shape, (1, 3))
 
+  def test_what_if_you_lower_fun_something_with_internal_effects(self):
+    bjp_p = core.Primitive('bjp')
+
+    @bjp_p.def_abstract_eval
+    def _(aval):
+      return aval
+
+    def lowering(x):
+      x_ref = jax.new_ref(x)
+      x_ref[...] += 1
+      x_ref[...] += -1
+      return jax.freeze(x_ref)
+
+    mlir.register_lowering(bjp_p, mlir.lower_fun(lowering, multiple_results=False))
+
+    @jax.jit
+    def f(x):
+      return bjp_p.bind(x)
+
+    f(3.)  # don't crash
+
+
 @jtu.with_config(jax_mutable_array_checks=True)
 class MutableArrayErrorsTest(jtu.JaxTestCase):
   def test_return_from_jit(self):