fix box bug in scan transpose

Author: mattjj

jax/_src/lax/control_flow/loops.py

@@ -876,14 +876,15 @@ def _scan_transpose(cts, *args, reverse, length, num_consts,
   jaxpr_trans = pe.move_outvars_to_back(
       jaxpr_trans, appends_out + [False] * (len(jaxpr_trans.out_avals) - len(appends_out)))
   num_attr_carry = sum(init_tree.num_leaves for init_tree, _, (_, _, kind)
-                       in attrs_tracked if kind is pe.ReadWrite)
+                       in attrs_tracked if kind in (pe.ReadWrite, pe.BoxAttr))
   linear_trans = ([False] * num_ires + [False] * num_attr_carry +
                   [True] * (len(ct_consts) + len(ct_carry) + len(ct_ys)) +
                   [False] * num_eres)
   in_state = _get_states(attrs_tracked)
 
   transpose_inputs = *ires, *in_state, *ct_consts, *ct_carry, *ct_ys, *eres
   transpose_num_out_carry = num_consts-num_ires+num_carry+num_attr_carry
+  assert len(transpose_inputs) == len(linear_trans), breakpoint()
 
   if not _split_transpose:
     outs = scan_p.bind(

tests/attrs_test.py

@@ -16,6 +16,7 @@
 
 from dataclasses import dataclass
 import itertools as it
+from functools import partial
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -1325,6 +1326,50 @@ def f(lst1, lst2):
     with self.assertRaisesRegex(ValueError, "a List instance can't be passed"):
       f(b, b)
 
+  def test_scan_transpose_regression(self):
+    @partial(jax.custom_vjp, nondiff_argnums=(0,))
+    def g2(box, x, i):
+      return x
+
+    def g2_fwd(box, x, i):
+      return x, i
+
+    def g2_bwd(box, i, g):
+      boxg = g.reshape((1, *g.shape))
+      box.set(jax.lax.dynamic_update_slice(box.get(), boxg, (i, 0, 0)))
+      return g, None
+
+    g2.defvjp(g2_fwd, g2_bwd)
+
+    def block_box(x, w, box, index):
+      x = x @ w
+      x = g2(box, x, index)
+      return x
+
+    def fwd_box(ws, x, box, stack):
+      def internal_scan(carry, w):
+        x, index = carry
+        x = block_box(x, w, box, index)
+        return (x, index+1), None
+      
+      (x, _), _ = jax.lax.scan(internal_scan, (x, 0), ws)
+      return x
+
+    D = 16
+    L = 4
+    B = 2
+    x = jnp.ones((B, D))
+    y = jnp.ones((B, D))
+    ws = jnp.ones((L, D, D))
+    grad_box = Box(jnp.zeros((L, B, D)))
+    grad_list = {'block_out': List()}
+
+    def jax_loss(ws, inputs, targets, box_or_list):
+      preds = fwd_box(ws, inputs, box_or_list, stack=stack)
+      return jnp.square(preds - targets).mean()
+
+    attr_grads = jax.grad(jax_loss)(ws, x, y, grad_box)  # don't crash
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())