fix basic scan bug with attrs

jax/_src/lax/control_flow/loops.py

@@ -268,7 +268,8 @@ def _create_jaxpr(init):
  if len(out_tree_children) != 2:
  msg = "scan body output must be a pair, got {}."
  raise TypeError(msg.format(tree_unflatten(out_tree, jaxpr.out_avals)))
- carry_avals_out = jaxpr.out_avals[:out_tree_children[0].num_leaves]
+ _, carry_avals_out, _ = split_list(
+ jaxpr.out_avals, [len(attrs_tracked), out_tree_children[0].num_leaves])
  return (init_flat, carry_avals, carry_avals_out, init_tree, in_flat, jaxpr,
  consts, out_tree, out_tree_children, attrs_tracked)
 

tests/attrs_test.py

@@ -344,6 +344,21 @@ def jitted():
 
  jax.jit(jitted)() # don't crash
 
+ def test_scan_carry(self):
+ class A:
+ ...
+
+ a = A()
+
+ jax_setattr(a, 'x', jnp.zeros(3))
+
+ def body(i, _):
+ x = jax_getattr(a, 'x')
+ x = x.at[i].set(x[i] + 1)
+ jax_setattr(a, 'x', x)
+ return i + 1, None
+ _, _ = jax.lax.scan(body, 0, None, length=3) # don't crash
+
 
 class AttrsJVPTest(jtu.JaxTestCase):