PR jax-ml#27576: [attrs] experimental appendattr

Imported from GitHub PR jax-ml#27576

This is an experimental extension to attrs. Attrs should be considered both experimental and deprecated.

This PR also includes some fixes for getattr/setattr.
Copybara import of the project:

--
3b1ea1a by Matthew Johnson <[email protected]>:

[attrs] experimental appendattr

Merging this change closes jax-ml#27576

COPYBARA_INTEGRATE_REVIEW=jax-ml#27576 from mattjj:appendattr b937952
PiperOrigin-RevId: 741662724

Author: mattjj

jax/_src/interpreters/partial_eval.py

@@ -58,6 +58,13 @@ def identity(x): return x
 AvalId = int
 ConstId = int
 
+AttrKind = Any
+PyTree = Any
+
+# Attrs flavors, see jax/experimental/attrs.py
+ReadWrite = type('ReadWrite', (), {})()
+Append = type('Append', (), {})()
+
 def _update_annotation_known(
     f: lu.WrappedFun,
     orig_type: InputType | None,
@@ -1553,6 +1560,17 @@ def move_binders_to_back(closed_jaxpr: ClosedJaxpr, to_move: Sequence[bool]
   """Reorder `invars` by moving those indicated in `to_move` to the back."""
   return move_binders_to_front(closed_jaxpr, map(op.not_, to_move))
 
+def move_outvars_to_back(jaxpr: ClosedJaxpr, to_move: Sequence[bool]) -> ClosedJaxpr:
+  return _move_outvars_to_back(jaxpr, tuple(to_move))
+
+@weakref_lru_cache
+def _move_outvars_to_back(jaxpr, to_move):
+  new_outvars = ([e for e, m in zip(jaxpr.jaxpr.outvars, to_move) if not m] +
+                 [e for e, m in zip(jaxpr.jaxpr.outvars, to_move) if     m])
+  return jaxpr.replace(jaxpr=jaxpr.jaxpr.replace(outvars=new_outvars))
+
+
+
 class DynamicJaxprTracer(core.Tracer):
   __slots__ = ['aval', '_debug_info']
 
@@ -1657,7 +1675,7 @@ class JaxprStackFrame:
   eqns: list[JaxprEqn]
   invars: list[Var]
   effects: core.Effects
-  attrs_tracked: list[tuple[Any, str]]
+  attrs_tracked: list[tuple[Any, str, AttrKind]]
   attrs_inits: list
   attrs_vars: list[Var]
   debug_info: core.DebugInfo
@@ -1679,10 +1697,14 @@ def __init__(self, debug_info: core.DebugInfo):
   def add_eqn(self, eqn: core.JaxprEqn):
     self.eqns.append(eqn)
 
-  def to_jaxpr(self, trace: DynamicJaxprTrace,
-               out_tracers: Sequence[Tracer],
-               debug_info: core.DebugInfo,
-               ) -> tuple[Jaxpr, list[Any], list[tuple[PyTreeDef, PyTreeDef, tuple[Any, str]]]]:
+  def reset_states(self):
+    reset_states(self.attrs_tracked, self.attrs_inits)
+
+  def to_jaxpr(
+      self, trace: DynamicJaxprTrace,
+      out_tracers: Sequence[Tracer],
+      debug_info: core.DebugInfo,
+    ) -> tuple[Jaxpr, list[Any], list[tuple[PyTreeDef, PyTreeDef, tuple[Any, str, AttrKind]]]]:
     # It's not necessary, but we keep the tracer-to-var mapping injective:
     assert len(self.tracer_to_var) == len(set(self.tracer_to_var.values()))
     invars = self.attrs_vars + self.invars
@@ -1699,7 +1721,6 @@ def to_jaxpr(self, trace: DynamicJaxprTrace,
     jaxpr, constvals = _const_folding_and_forwarding(jaxpr, constvals)
     jaxpr, constvals = _inline_literals(jaxpr, constvals)
     init_trees = [tree_structure(init_val) for init_val in self.attrs_inits]
-    set_states(self.attrs_tracked, self.attrs_inits)  # reset to initial values
     return jaxpr, list(constvals), zip(init_trees, end_trees, self.attrs_tracked)
 
   def to_jaxpr2(self, out_tracers: Sequence[core.Tracer],
@@ -1840,10 +1861,9 @@ def vars_in_shape(aval: AbstractValue) -> Sequence[Var]:
     outvars = [var(v) if v in used else dropvar(v.aval) for v in eqn.outvars]
     new_eqns.append(eqn.replace(invars=invars, outvars=outvars))
   new_outvars = [lit_or_var(v) for v in jaxpr.outvars]
-  jaxpr_effects = make_jaxpr_effects(new_constvars, new_invars, new_outvars,
-                                      new_eqns)
-  new_jaxpr = Jaxpr(new_constvars, new_invars, new_outvars, new_eqns,
-                    jaxpr_effects, jaxpr.debug_info)
+  effs = make_jaxpr_effects(new_constvars, new_invars, new_outvars, new_eqns)
+  new_jaxpr = Jaxpr(new_constvars, new_invars, new_outvars, new_eqns, effs,
+                    jaxpr.debug_info)
   return new_jaxpr, new_constvals
 
 
@@ -2172,19 +2192,23 @@ def trace_to_jaxpr_dynamic(
     *,
     keep_inputs: list[bool] | None = None,
 ) -> tuple[Jaxpr, list[AbstractValue], list[Any],
-           list[tuple[PyTreeDef, PyTreeDef, tuple[Any, str]]]]:
+           list[tuple[PyTreeDef, PyTreeDef, tuple[Any, str, AttrKind]]]]:
   keep_inputs = [True] * len(in_avals) if keep_inputs is None else keep_inputs
   trace = DynamicJaxprTrace(fun.debug_info)
   with core.ensure_no_leaks(trace), source_info_util.reset_name_stack():
     in_tracers = _input_type_to_tracers(trace.new_arg, in_avals)
     in_tracers = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-    with core.set_current_trace(trace):
-      ans = fun.call_wrapped(*in_tracers)
+    try:
+      with core.set_current_trace(trace):
+        ans = fun.call_wrapped(*in_tracers)
 
-    out_tracers = map(trace.to_jaxpr_tracer, ans)
-    _check_no_returned_refs(fun.debug_info, out_tracers)
-    jaxpr, consts, attrs_tracked = trace.to_jaxpr(out_tracers, fun.debug_info)
-    del trace, fun, in_tracers, out_tracers, ans
+      out_tracers = map(trace.to_jaxpr_tracer, ans)
+      _check_no_returned_refs(fun.debug_info, out_tracers)
+      jaxpr, consts, attrs_tracked = trace.to_jaxpr(out_tracers, fun.debug_info)
+      del fun, in_tracers, out_tracers, ans
+    finally:
+      trace.frame.reset_states()
+      del trace
 
   config.enable_checks.value and core.check_jaxpr(jaxpr)
   return jaxpr, [v.aval for v in jaxpr.outvars], consts, attrs_tracked
@@ -2242,14 +2266,14 @@ def trace_to_jaxpr_dynamic2(
     tuple[AbstractedAxisName, ...],
 ]
 
-AttrsTracked = list[tuple[Any, str]]
+AttrsTracked = list[tuple[Any, str, AttrKind]]
 AttrStates = list
-def set_states(attrs_tracked: AttrsTracked, vals: AttrStates):
-  for ((obj, attr), val) in zip(attrs_tracked, vals):
+def reset_states(attrs_tracked: AttrsTracked, init_vals: AttrStates) -> None:
+  for ((obj, attr, _), val) in zip(attrs_tracked, init_vals):
     setattr(obj, attr, val) if val is not dne_sentinel else delattr(obj, attr)
 
-def get_states(attrs_tracked: AttrsTracked):
-  return [getattr(obj, attr) for (obj, attr) in attrs_tracked]
+def get_states(attrs_tracked: AttrsTracked) -> list[PyTree]:
+  return [getattr(obj, attr) for (obj, attr, kind) in attrs_tracked]
 
 @register_static
 class DoesNotExist: ...

jax/_src/lax/control_flow/loops.py

@@ -298,8 +298,17 @@ 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, _ = split_list(
-        jaxpr.out_avals, [len(attrs_tracked), out_tree_children[0].num_leaves])
+
+    if attrs_tracked:
+      appends_out = [kind is pe.Append for *_, (_, _, kind) in attrs_tracked]
+      jaxpr = pe.move_outvars_to_back(
+          jaxpr, appends_out + [False] * (len(jaxpr.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)
+      _, carry_avals_out, _ = split_list(
+          jaxpr.out_avals, [num_attr_carry, out_tree_children[0].num_leaves])
+    else:
+      carry_avals_out, _ = split_list(jaxpr.out_avals, [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)
 
@@ -332,37 +341,59 @@ def _create_jaxpr(init):
     raise ValueError("`unroll` must be a `bool` or a positive `int`.")
   if attrs_tracked:
     in_state = _get_states(attrs_tracked)
-    in_carry, in_ext = split_list(in_flat, [num_carry])
-    in_flat = [*in_state, *in_carry, *in_ext]
-    num_carry += len(attrs_tracked)
+    in_flat = [*in_state, *in_flat]
+    num_carry += len(in_state)
   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)
   if attrs_tracked:
-    out_state, out = split_list(out, [len(attrs_tracked)])
-    _set_states(attrs_tracked, out_state)
+    num_ext = (len(out) - len(in_state)
+               - sum(k is pe.Append for *_, (_, _, k) in attrs_tracked))
+    out_state, out, out_append = split_list(out, [len(in_state), num_ext])
+    out_attrs = _merge_attrs_out(attrs_tracked, out_state, out_append)
+    _set_states(attrs_tracked, out_attrs)
   return tree_unflatten(out_tree, out)
 
 def _set_states(attrs_tracked, vals):
-  from jax.experimental.attrs import jax_setattr
+  from jax.experimental.attrs import jax_setattr, jax_extendattr
   valss = split_list_checked(vals, [td.num_leaves for _, td, _ in attrs_tracked])
-  for ((_, treedef, (obj, attr)), leaves) in zip(attrs_tracked, valss):
-    val = tree_unflatten(treedef, leaves)
-    jax_setattr(obj, attr, val)
+  for ((_, treedef, (obj, attr, kind)), leaves) in zip(attrs_tracked, valss):
+    if kind is pe.ReadWrite:
+      val = tree_unflatten(treedef, leaves)
+      jax_setattr(obj, attr, val)
+    elif kind is pe.Append:
+      val, = leaves
+      jax_extendattr(obj, attr, val.reshape(-1, *val.shape[2:]))
+    else:
+      assert False
 
 def _get_states(attrs_tracked):
   from jax.experimental.attrs import jax_getattr
   vals = []
-  for treedef, _, (obj, attr) in attrs_tracked:
-    tree = jax_getattr(obj, attr)
-    leaves, treedef_ = tree_flatten(tree)
-    assert treedef == treedef_
-    vals.extend(leaves)
+  for treedef, _, (obj, attr, kind) in attrs_tracked:
+    if kind is pe.ReadWrite:
+      tree = jax_getattr(obj, attr)
+      leaves, treedef_ = tree_flatten(tree)
+      assert treedef == treedef_
+      vals.extend(leaves)
+    elif kind is pe.Append:
+      pass
+    else:
+      assert False
   return vals
 
+def _merge_attrs_out(attrs_tracked, out_state, out_append):
+  out_state_, out_append_ = iter(out_state), iter(out_append)
+  out_attrs = [item for _, out_tree, (_, _, k) in attrs_tracked for item in
+               (itertools.islice(out_state_, out_tree.num_leaves)
+               if k is pe.ReadWrite else [next(out_append_)])]
+  assert next(out_state_, None) is next(out_append_, None) is None
+  return out_attrs
+
+
 def _capitalize(s):
   # s.capitalize() converts s[1:] to lowercase which we don't want.
   return s[0].capitalize() + s[1:]
@@ -662,7 +693,7 @@ def _scan_partial_eval(trace, *tracers, reverse: bool,
   # The above trace_to_jaxpr_nounits call computed loop-invariant residuals
   # (known values in invar_pvals_out) and also computed loop-invariant values
   # needed by the new jaxpr_known (in jaxpr_known_consts, which replace the
-  # previous consts). We need to collect the computed inteisive residuals, and
+  # previous consts). We need to collect the computed intensive residuals, and
   # move corresponding intensive residual binders in jaxpr_unknown to the front.
   res_pvals = invar_pvals_out[len(invar_pvals_out) - num_res:]
   intensive_res = [pval.get_known() for pval in res_pvals if pval.is_known()]
@@ -785,16 +816,21 @@ def _scan_transpose(cts, *args, reverse, length, num_consts,
   ct_consts = _map(ad_util.zeros_like_aval, jaxpr.in_avals[num_ires:num_consts])
 
   #       jaxpr :: [ires, T d] -> [T c] -> [T a, eres] -> ([T c], [T b])
-  # jaxpr_trans :: [ires] -> [CT d, CT c] -> [CT b, eres] -> ([CT d, CT c], [CT a])
+  # jaxpr_trans :: [ires] -> [CT d, CT c] -> [CT b, eres] -> ([CT d, CT c], [CT a, e])
   jaxpr_trans, attrs_tracked = _transpose_scan_jaxpr(
       jaxpr, num_ires, num_consts - num_ires, num_eres, ct_ys_is_zeros)
-  linear_trans = ([False] * num_ires + [False] * len(attrs_tracked) +
+  appends_out = [kind is pe.Append for *_, (_, _, kind) in attrs_tracked]
+  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)
+  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+len(attrs_tracked)
+  transpose_num_out_carry = num_consts-num_ires+num_carry+num_attr_carry
 
   if not _split_transpose:
     outs = scan_p.bind(
@@ -889,8 +925,10 @@ def _scan_transpose(cts, *args, reverse, length, num_consts,
         for mask in outs_mask
     ]
 
-  out_state, outs = split_list(outs, [len(attrs_tracked)])
-  _set_states(attrs_tracked, out_state)
+  num_outs = len(outs) - num_attr_carry - sum(appends_out)
+  out_state, outs, out_append = split_list(outs, [num_attr_carry, num_outs])
+  out_attrs = _merge_attrs_out(attrs_tracked, out_state, out_append)
+  _set_states(attrs_tracked, out_attrs)
   ct_consts, ct_init, ct_xs = split_list(outs, [num_consts - num_ires, num_carry])
   return [None] * num_ires + ct_consts + ct_init + ct_xs + [None] * num_eres
 
@@ -935,12 +973,10 @@ def transposed(*res1_cbar_bbar_res2):
     return c_bar + a_bar
 
   # TODO(necula): fix arg names and results for transposed
-  transposed_wrapped = lu.wrap_init(transposed,
-                                    debug_info=jaxpr.jaxpr.debug_info)
-  return _make_closed_jaxpr_attrs(
-      transposed_wrapped,
-      tuple(res1_avals + c_avals + b_carry_avals +
-            b_ys_avals_stripped + res2_avals))
+  transposed_wrapped = lu.wrap_init(transposed, debug_info=jaxpr.jaxpr.debug_info)
+  trans_avals = (*res1_avals, *c_avals, *b_carry_avals, *b_ys_avals_stripped, *res2_avals)
+  trans_jaxpr, attrs_tracked = _make_closed_jaxpr_attrs(transposed_wrapped, trans_avals)
+  return trans_jaxpr, attrs_tracked
 
 
 def _scan_batching_rule(axis_data, args,