Apply extensive input to extensive output forwarding in scan.

Author: dfm

jax/_src/lax/control_flow/loops.py

@@ -341,7 +341,7 @@ def _create_jaxpr(init):
   # If the body forwards an input carry to an output carry, that input is
   # read-only and can be moved to be a const. Doing so can lead to efficiency
   # wins, e.g. if the scan is inside a cond with a batched predicate.
-  carry_fwd, _ = split_list(pe._jaxpr_forwarding(jaxpr.jaxpr), [num_carry])
+  carry_fwd, in_fwd = split_list(pe._jaxpr_forwarding(jaxpr.jaxpr), [num_carry])
   move_to_const = [len(consts) + i == f for i, f in enumerate(carry_fwd)]
   if any(move_to_const):
     jaxpr = pe.prune_closed_jaxpr_outputs(
@@ -352,12 +352,30 @@ def _create_jaxpr(init):
     consts = [*new_consts, *consts]
     num_carry -= len(new_consts)
 
+  # When an extensive output is forwarded from an extensive input, we can
+  # avoid copying it by pruning it from the jaxpr and forwarding manually. We
+  # don't need to update the indexing based on the optimization above since it
+  # doesn't change the total number of consts and carries combined. But, we do
+  # remove the number of consts from the index since `_jaxpr_forwarding`
+  # includes the Jaxpr's constvars when indexing.
+  in_fwd = [in_idx - len(consts) if in_idx is not None and
+            in_idx >= num_carry + len(consts) else None for in_idx in in_fwd]
+  jaxpr = pe.prune_closed_jaxpr_outputs(
+      jaxpr, [True] * num_carry + [i is None for i in in_fwd])
+
   out = scan_p.bind(*consts, *in_flat,
                     reverse=reverse, length=length, jaxpr=jaxpr,
                     num_consts=len(consts), num_carry=num_carry,
                     linear=(False,) * (len(consts) + len(in_flat)),
                     unroll=unroll, _split_transpose=_split_transpose)
 
+  # Apply input to output forwarding that was computed above.
+  consts_out, out = split_list(out, [num_carry])
+  out_ = iter(out)
+  out = [next(out_) if f is None else _maybe_put(in_flat[f]) for f in in_fwd]
+  assert next(out_, None) is None
+  out = [*consts_out, *out]
+
   if any(move_to_const):
     out = pe.merge_lists(move_to_const + [False] * num_ys, out, new_consts)
 

tests/lax_control_flow_test.py

@@ -3311,6 +3311,14 @@ def g(x):
     eqn_jaxpr = jaxpr.eqns[0].params["jaxpr"]
     self.assertIn("debug_callback", [e.primitive.name for e in eqn_jaxpr.eqns])
 
+  def test_scan_input_to_output_forwarding(self):
+    def f(c, x):
+      return c + 1, x
+    def g(x):
+      return jax.lax.scan(f, 0, x)
+    jaxpr = jax.make_jaxpr(g)(jnp.arange(3.))
+    self.assertLen(jaxpr.eqns[0].params["jaxpr"].jaxpr.outvars, 1)
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())