Merge pull request #32434 from mattjj:32399

Google-ML-Automation · Google-ML-Automation · commit a15ac6b66981 · 2025-10-07T17:58:07.000-07:00
PiperOrigin-RevId: 816456798
diff --git a/jax/_src/core.py b/jax/_src/core.py
@@ -3548,7 +3548,14 @@ def substitute(aval: AbstractValue):
   out_type = [a.update(shape=tuple(in_map.get(d, out_map.get(d))
                                    if type(d) is Var else d for d in a.shape))
               if type(a) is DShapedArray else a for a in out_avals]
-  return out_type, call_jaxpr.effects
+
+  # jaxpr input effects are indexed to include jaxpr.constvars, but the eqn
+  # should have effects indexed only on its explicit arguments
+  effs = {e.replace(input_index=e.input_index - len(call_jaxpr.constvars))
+          if isinstance(e, effects.JaxprInputEffect)
+          else e for e in call_jaxpr.effects}
+
+  return out_type, effs
 
 def _check_map(ctx_factory, prim, in_avals, params):
   if "call_jaxpr" not in params:
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
@@ -1980,7 +1980,15 @@ def _pjit_typecheck(ctx_factory, *in_atoms, jaxpr, **params):
 
 
 def _pjit_abstract_eval(*args, jaxpr, out_shardings, **_):
-  return jaxpr.out_avals, jaxpr.effects
+  # jaxpr input effects are indexed to include jaxpr.constvars, but the pjit eqn
+  # should have effects indexed only on its explicit arguments
+  if jaxpr.constvars:
+    effs = {e.replace(input_index=e.input_index - len(jaxpr.constvars))
+            if isinstance(e, effects.JaxprInputEffect)
+            else e for e in jaxpr.effects}
+  else:
+    effs = jaxpr.effects
+  return jaxpr.out_avals, effs
 jit_p.def_effectful_abstract_eval(_pjit_abstract_eval)
 
 
diff --git a/tests/jaxpr_effects_test.py b/tests/jaxpr_effects_test.py
@@ -257,6 +257,26 @@ def f(x):
       jaxpr = jax.make_jaxpr(f)(np.arange(jax.local_device_count()))
     self.assertSetEqual(jaxpr.effects, {foo_effect, bar_effect})
 
+  def test_pjit_const_input_effect_indexing(self):
+    # https://github.com/jax-ml/jax/issues/32399
+    @jax.jit
+    def bar(x, w):
+        def scan_fn(x, _):
+            c = jnp.array([])
+            o = w[...] @ x
+            x = jnp.concatenate([x, c], axis=-1)
+            return x, None
+
+        x, _ = jax.lax.scan(scan_fn, x, None, length=10)
+        return x
+
+
+    @jax.jit
+    def foo(w):
+        return bar(jnp.zeros((1,)), w)
+
+    foo(jax.new_ref(jnp.eye(1)))  # don't crash
+
 
 @jtu.thread_unsafe_test_class()  # because of mlir.register_lowering calls
 class EffectfulJaxprLoweringTest(jtu.JaxTestCase):
