Merge pull request #20278 from mattjj:fix-forwarding-3

PiperOrigin-RevId: 616264716

Author: jax authors

jax/_src/interpreters/partial_eval.py

@@ -847,7 +847,7 @@ def trace_to_subjaxpr_nounits_fwd2(
   out_pvals = [t.pval for t in out_tracers]
 
   # Which consts (aka residuals) are just forwarded inputs? Check obj id.
-  in_consts  = [pval.get_known()    for pval in  in_pvals if     pval.is_known()]
+  in_consts  = [pval.get_known()    for pval in  in_pvals if    pval.is_known()]
   id_map = {id(c): i for i, c in enumerate(in_consts)}
   input_fwds: list[int | None] = [id_map.get(id(c)) for c in consts]
 

tests/api_test.py

@@ -4593,14 +4593,16 @@ def inner(a, x):
   def test_jit_forwarding_correctness(self, seed, num_input_fwd, num_output_fwd):
     num_args = 3
     rng = np.random.RandomState(seed)
+    in_perm = rng.permutation(num_args)
+    out_perm = rng.permutation(num_args)
 
     @jax.jit
     def f(inputs):
-      inputs = [inputs[i] for i in rng.permutation(num_args)]
+      inputs = [inputs[i] for i in in_perm]
       outputs = inputs[:num_input_fwd] + [
           jnp.exp(inputs[i]) if i < num_output_fwd else jnp.sin(inputs[i])
           for i in range(num_args - num_input_fwd)]
-      return [outputs[i] for i in rng.permutation(num_args)]
+      return [outputs[i] for i in out_perm]
 
     jtu.check_grads(f, (list(jnp.arange(float(num_args))),), order=1,
                     modes=['rev'], atol=1e-3, rtol=1e-3)

tests/shard_map_test.py

@@ -1604,6 +1604,27 @@ def f(x):
     with jax.disable_jit():
       f(x)  # don't crash
 
+  @parameterized.parameters(it.product(range(4), repeat=3))
+  @jtu.run_on_devices("cpu")
+  def test_forwarding_correctness(self, seed, num_input_fwd, num_output_fwd):
+    num_args = 3
+    rng = np.random.RandomState(seed)
+    mesh = Mesh(np.array(jax.devices()[:1]), ('i',))
+
+    in_perm = rng.permutation(num_args)
+    out_perm = rng.permutation(num_args)
+
+    @partial(shard_map, mesh=mesh, in_specs=P('i'), out_specs=P('i'))
+    def f(inputs):
+      inputs = [inputs[i] for i in in_perm]
+      outputs = inputs[:num_input_fwd] + [
+          jnp.exp(inputs[i]) if i < num_output_fwd else jnp.sin(inputs[i])
+          for i in range(num_args - num_input_fwd)]
+      return [outputs[i] for i in out_perm]
+
+    jtu.check_grads(f, (list(jnp.arange(float(num_args))[:,None]),), order=1,
+                    modes=['rev'], atol=1e-3, rtol=1e-3)
+
 
 class FunSpec(NamedTuple):
   name: str