Repository: google/jax

Issue:
`grad(jit(f))` can fail when an output which is a forwarded input precedes an output which is also a residual
### Description

Example:

```python
import jax
import jax.numpy as jnp

def f(x):
    @jax.jit
    def inner(a, x):
        return a, jnp.exp(x)
    return inner(0., x)[0]
jax.grad(f)(1.)
# IndexError: list index out of range
```
The error occurs in the call to `subs_list2` [here](https://github.com/google/jax/blob/a5d32c41c336368d8ec01bda17072d34189f7faa/jax/_src/pjit.py#L1844). It looks like the problem is that `idx_map` is built using indices into `out_vars`, so initially the `exp(x)` value is marked as being in output index 1. When the forwarded inputs and duplicated residuals are pruned, this index doesn't get updated, and can end up being out of bounds.

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.25
jaxlib: 0.4.25
numpy:  1.25.2
python: 3.11.6 (main, Nov 14 2023, 09:36:21) [GCC 13.2.1 20230801]
jax.devices (2 total, 2 local): [cuda(id=0) cuda(id=1)]
process_count: 1
```


Discussion:
Oof, what a bad bug! Thanks for reporting this!

Resolving patch:
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -76,7 +76,7 @@
 from jax._src.util import (
     HashableFunction, safe_map, safe_zip, wraps,
     distributed_debug_log, split_list, weakref_lru_cache,
-    merge_lists, flatten, unflatten, subs_list2)
+    merge_lists, flatten, unflatten)
 
 map, unsafe_map = safe_map, map
 zip, unsafe_zip = safe_zip, zip
@@ -1798,8 +1798,8 @@ def _pjit_partial_eval(trace, *in_tracers,
 
   known_ins = tuple(pv.is_known() for pv in in_pvals)
   unknown_ins = tuple(not k for k in known_ins)
-  known_jaxpr, unknown_jaxpr, unknown_outs, res_avals = pe.partial_eval_jaxpr_nounits(
-      jaxpr, unknown_ins, instantiate=False)
+  known_jaxpr, unknown_jaxpr, unknown_outs, res_avals = \
+      pe.partial_eval_jaxpr_nounits(jaxpr, unknown_ins, instantiate=False)
   unknown_outs = tuple(unknown_outs)
   known_outs = tuple(not uk for uk in unknown_outs)
   num_residuals = len(res_avals)
@@ -1808,28 +1808,37 @@ def _pjit_partial_eval(trace, *in_tracers,
   def keep_where(l, should_keep):
     return tuple(x for x, keep in zip(l, should_keep) if keep)
 
-  # Compute which outputs are just forwarded inputs.
-  num_out_primals = len(known_jaxpr.out_avals) - num_residuals
-  in_fwd = pe._jaxpr_forwarding(known_jaxpr.jaxpr)
-
-  # Only forward primal outputs when corresponding out_sharding is UNSPECIFIED.
-  in_fwd_primal, in_fwd_res = split_list(in_fwd, [num_out_primals])
-  in_fwd = [fwd if is_unspecified(os) else None for os, fwd in
-            zip(keep_where(out_shardings, known_outs), in_fwd_primal)
-            ] + in_fwd_res
-  del in_fwd_primal, in_fwd_res
-
-  # Compute which residuals are just primal outputs.
-  out_vars, res_vars = split_list(known_jaxpr.jaxpr.outvars, [num_out_primals])
-  idx_map = {id(v): i for i, v in enumerate(out_vars)}
-  out_fwd = [None] * num_out_primals + [idx_map.get(id(v)) for v in res_vars]
-
-  # Prune jaxpr outputs and out_shardings by removing forwards.
-  keep = [f1 is None and f2 is None for f1, f2 in zip(in_fwd, out_fwd)]
-  known_jaxpr = pe.prune_closed_jaxpr_outputs(known_jaxpr, keep)
   known_out_shardings = keep_where(out_shardings, known_outs) + res_shardings
-  known_out_shardings = keep_where(known_out_shardings, keep)
-  del keep, num_out_primals
+
+  # TODO(mattjj): un-disable this optimization after we have more tests
+  # # Input-to-output forwarding: compute which outputs are just forwarded inputs.
+  # num_out_primals = len(known_jaxpr.out_avals) - num_residuals
+  # in_fwd: list[int | None] = pe._jaxpr_forwarding(known_jaxpr.jaxpr)
+  # # Only forward primal outputs when corresponding out_sharding is UNSPECIFIED.
+  # in_fwd_primal, in_fwd_res = split_list(in_fwd, [num_out_primals])
+  # in_fwd = [fwd if is_unspecified(os) else None for os, fwd in
+  #           zip(keep_where(out_shardings, known_outs), in_fwd_primal)
+  #           ] + in_fwd_res
+  # del in_fwd_primal, in_fwd_res
+  # # Prune jaxpr outputs and out_shardings by removing the input-forwards.
+  # keep = [f is None for f in in_fwd]
+  # known_jaxpr = pe.prune_closed_jaxpr_outputs(known_jaxpr, keep)
+  # known_out_shardings = keep_where(known_out_shardings, keep)
+  # # Update num_out_primals to reflect pruning.
+  # kept_primals, kept_res = split_list(keep, [num_out_primals])
+  # num_out_primals = sum(f is None for f in kept_primals)
+  # del keep, kept_primals, kept_res
+
+  # TODO(mattjj): un-disable this optimization after we have more tests
+  # # Output-to-output forwarding: compute which residuals are just primal outputs
+  # out_vars, res_vars = split_list(known_jaxpr.jaxpr.outvars, [num_out_primals])
+  # idx_map = {id(v): i for i, v in enumerate(out_vars)}
+  # out_fwd = [None] * num_out_primals + [idx_map.get(id(v)) for v in res_vars]
+  # # Prune jaxpr outputs and out_shardings by removing forwarded residuals.
+  # keep = [f is None for f in out_fwd]
+  # known_jaxpr = pe.prune_closed_jaxpr_outputs(known_jaxpr, keep)
+  # known_out_shardings = keep_where(known_out_shardings, keep)
+  # del keep
 
   known_params = dict(
       jaxpr=known_jaxpr, in_shardings=keep_where(in_shardings, known_ins),
@@ -1841,16 +1850,19 @@ def keep_where(l, should_keep):
   # Bind known things to pjit_p.
   known_inputs = [pv.get_known() for pv in in_pvals if pv.is_known()]
   all_known_outs = pjit_p.bind(*known_inputs, **known_params)
-  all_known_outs = subs_list2(in_fwd, out_fwd, known_inputs, all_known_outs,
-                              all_known_outs)
+  # TODO(mattjj): un-disable this optimization after we have more tests
+  # # Add back in the output fwds.
+  # all_known_outs = subs_list(out_fwd, all_known_outs, all_known_outs)
+  # # Add back in the input fwds.
+  # all_known_outs = subs_list(in_fwd, known_inputs, all_known_outs)
 
   known_out_vals, residual_vals = \
       split_list(all_known_outs, [len(all_known_outs) - num_residuals])
   residual_tracers = map(trace.new_instantiated_const, residual_vals)
 
-  # The convention of partial_eval_jaxpr_nounits is to place residual binders
-  # at the front of the jaxpr produced, so we move them to the back since both
-  # the jaxpr equation built below and the pjit transpose rule assume a
+  # The convention of partial_eval_jaxpr_nounits is to place residual binders at
+  # the front of the jaxpr produced, so we move them to the back since both the
+  # jaxpr equation built below and the pjit transpose rule assume a
   # residual-inputs-last convention.
   unknown_jaxpr = pe.move_binders_to_back(
       unknown_jaxpr, [True] * num_residuals + [False] * sum(unknown_ins))


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4565,6 +4565,15 @@ def foo(self):
     gc.collect()
     assert a() is None
 
+  def test_forwarding_bug(self):
+    # Test for issue #20267.
+    def f(x):
+        @jax.jit
+        def inner(a, x):
+            return a, jnp.exp(x)
+        return inner(0., x)[0]
+    jax.grad(f)(1.)  # don't crash
+
 
 class RematTest(jtu.JaxTestCase):
 
diff --git a/tests/core_test.py b/tests/core_test.py
--- a/tests/core_test.py
+++ b/tests/core_test.py
@@ -368,6 +368,8 @@ def body(c, _):
     dropvar, b = jaxpr.eqns[0].outvars
     self.assertEqual(dropvar.aval, aval)
 
+  # TODO(mattjj): un-skip
+  @unittest.skip('temporarily skipping until we can add more tests')
   def test_input_residual_forwarding(self):
     # https://github.com/google/jax/pull/11151
     x = jnp.arange(3 * 4.).reshape(3, 4)