braintrace 0.2.5__tar.gz → 0.2.6__tar.gz

{braintrace-0.2.5 → braintrace-0.2.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.5
+Version: 0.2.6
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_algorithm/vjp_graph_executor.py

@@ -49,7 +49,7 @@ from brainstate._compatible_import import get_aval
 from jax.interpreters import partial_eval as pe
 from jax.tree_util import register_pytree_node_class
 
-from braintrace._compatible_imports import Var, wrap_init
+from braintrace._compatible_imports import Var, jaxpr_constvars, wrap_init
 from braintrace._compiler import (
     ControlFlowPolicy,
     DEFAULT_MAX_JACOBIAN_ELEMENTS,
@@ -341,7 +341,7 @@ class ETraceVjpGraphExecutor(ETraceGraphExecutor):
                 y_stack = intermediate_values[m[relation.y_var]]
                 cvars = list(dict.fromkeys(
                     v for j in relation.y_to_hidden_group_jaxprs
-                    for v in j.constvars))
+                    for v in jaxpr_constvars(j, 1)))
                 c_stacks = tuple(intermediate_values[m[v]] for v in cvars)
 
                 def _one_substep(

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compatible_imports.py

@@ -32,9 +32,12 @@ __all__ = [
     'scan_num_consts_carry',
     'scan_params_add_ys',
     'wrap_init',
+    'jaxpr_all_invars',
+    'split_jaxpr_invars',
+    'jaxpr_constvars',
 ]
 
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, List, Tuple
 
 from brainstate._compatible_import import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal, wrap_init
 
@@ -146,3 +149,132 @@ def scan_params_add_ys(params: Dict, n_extra: int) -> Dict:
     carry_ft, ys_ft = params['ft_out'].unpack()
     new_ft_out = _ft.pack((carry_ft, _ft.pack((ys_ft, _ft.nones(n_extra)))))
     return {**params, 'ft_out': new_ft_out}
+
+
+def jaxpr_all_invars(jaxpr: Jaxpr) -> List[Var]:
+    """Return every positional input of *jaxpr*, constvars first.
+
+    This is the argument order :func:`jax.core.eval_jaxpr` expects: it writes
+    ``[*consts, *args]`` against exactly this list.
+
+    JAX 0.11 merged ``ClosedJaxpr`` into ``Jaxpr`` and now derives the
+    ``constvars`` / ``invars`` boundary from how many const *values* are
+    attached, rather than storing it with the symbols. An open jaxpr built with
+    symbolic constvars and no attached values therefore reports
+    ``constvars == []`` and ``invars == [*constvars, *invars]``. Their
+    concatenation is stable across every supported JAX version, which is why
+    every split in this module is derived from it.
+
+    Parameters
+    ----------
+    jaxpr : Jaxpr
+        The jaxpr to inspect.
+
+    Returns
+    -------
+    list of Var
+        The concatenation ``[*jaxpr.constvars, *jaxpr.invars]``.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> from braintrace._compatible_imports import jaxpr_all_invars
+        >>> jaxpr = jax.make_jaxpr(lambda x, y: x * y)(jnp.ones(3), jnp.ones(3))
+        >>> len(jaxpr_all_invars(jaxpr.jaxpr))
+        2
+    """
+    return [*jaxpr.constvars, *jaxpr.invars]
+
+
+def split_jaxpr_invars(jaxpr: Jaxpr, num_invars: int) -> Tuple[List[Var], List[Var]]:
+    """Split *jaxpr*'s positional inputs into ``(constvars, invars)``.
+
+    The compiler builds transition jaxprs whose ``invars`` are the
+    differentiated inputs and whose ``constvars`` are surrounding intermediates
+    bound at execution time. Since JAX 0.11 that boundary cannot be read back
+    off the jaxpr (see :func:`jaxpr_all_invars`), so it is recovered from the
+    invar count, which every caller already knows -- the invars are what it is
+    about to feed in.
+
+    Parameters
+    ----------
+    jaxpr : Jaxpr
+        The jaxpr to split.
+    num_invars : int
+        The number of trailing positional inputs that are ``invars``. The
+        leading remainder are the ``constvars``.
+
+    Returns
+    -------
+    tuple of (list of Var, list of Var)
+        ``(constvars, invars)``.
+
+    Raises
+    ------
+    ValueError
+        If *num_invars* is negative or exceeds the number of positional inputs.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> from braintrace._compatible_imports import split_jaxpr_invars
+        >>> jaxpr = jax.make_jaxpr(lambda x, y: x * y)(jnp.ones(3), jnp.ones(3))
+        >>> constvars, invars = split_jaxpr_invars(jaxpr.jaxpr, 1)
+        >>> len(constvars), len(invars)
+        (1, 1)
+    """
+    all_invars = jaxpr_all_invars(jaxpr)
+    if not 0 <= num_invars <= len(all_invars):
+        raise ValueError(
+            f'Cannot split a jaxpr with {len(all_invars)} positional inputs at '
+            f'num_invars={num_invars}. This means the caller and the jaxpr '
+            f'disagree about the transition signature; evaluating it would '
+            f'misalign the eval_jaxpr argument list.'
+        )
+    boundary = len(all_invars) - num_invars
+    return all_invars[:boundary], all_invars[boundary:]
+
+
+def jaxpr_constvars(jaxpr: Jaxpr, num_invars: int) -> List[Var]:
+    """Return only the ``constvars`` of *jaxpr*, given its invar count.
+
+    Thin wrapper over :func:`split_jaxpr_invars` for the common case where the
+    caller needs the const values to pass to :func:`jax.core.eval_jaxpr` and
+    already holds the invars.
+
+    Parameters
+    ----------
+    jaxpr : Jaxpr
+        The jaxpr to inspect.
+    num_invars : int
+        The number of trailing positional inputs that are ``invars``.
+
+    Returns
+    -------
+    list of Var
+        The leading ``constvars``.
+
+    Raises
+    ------
+    ValueError
+        If *num_invars* is negative or exceeds the number of positional inputs.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> from braintrace._compatible_imports import jaxpr_constvars
+        >>> jaxpr = jax.make_jaxpr(lambda x, y: x * y)(jnp.ones(3), jnp.ones(3))
+        >>> len(jaxpr_constvars(jaxpr.jaxpr, 1))
+        1
+    """
+    constvars, _ = split_jaxpr_invars(jaxpr, num_invars)
+    return constvars

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compatible_imports_test.py

@@ -15,12 +15,16 @@
 
 import types
 
+import jax
 import jax.numpy as jnp
+import pytest
 from jax import jit, make_jaxpr, lax
 
 from braintrace._compatible_imports import (
+    Jaxpr,
     is_jit_primitive, is_scan_primitive, is_while_primitive,
     is_cond_primitive, scan_num_consts_carry, scan_params_add_ys,
+    jaxpr_all_invars, split_jaxpr_invars, jaxpr_constvars,
 )
 
 
@@ -81,6 +85,76 @@ class TestScanAddYs:
         assert len(new_ys) == len(old_ys) + 2        # ys grew by 2
 
 
+def _open_jaxpr_with_split(num_constvars: int):
+    """An open jaxpr built the way the compiler builds transition jaxprs.
+
+    Symbolic ``constvars`` with no attached const *values* -- the shape JAX
+    0.11 no longer records a boundary for. Returns ``(jaxpr, constvars,
+    invars)`` so a test can compare against the intended split rather than
+    whatever the object reports.
+    """
+
+    def f(a, b, h):
+        return a * h + b
+
+    traced = make_jaxpr(f)(jnp.ones(3), jnp.ones(3), jnp.ones(3)).jaxpr
+    all_in = jaxpr_all_invars(traced)
+    constvars, invars = all_in[:num_constvars], all_in[num_constvars:]
+    jaxpr = Jaxpr(
+        constvars=constvars,
+        invars=invars,
+        outvars=traced.outvars,
+        eqns=traced.eqns,
+        debug_info=traced.debug_info,
+    )
+    return jaxpr, constvars, invars
+
+
+class TestJaxprInvarSplit:
+    """The const/invar split of a value-less open jaxpr survives JAX 0.11.
+
+    JAX 0.11 merged ``ClosedJaxpr`` into ``Jaxpr`` and derives
+    ``constvars``/``invars`` from how many const *values* are attached, so an
+    open jaxpr built with symbolic constvars reports ``constvars == []``. These
+    helpers recover the boundary from the known invar count instead.
+    """
+
+    def test_all_invars_is_the_eval_jaxpr_argument_order(self):
+        jaxpr, constvars, invars = _open_jaxpr_with_split(2)
+        assert jaxpr_all_invars(jaxpr) == [*constvars, *invars]
+
+    def test_split_recovers_the_intended_boundary(self):
+        jaxpr, constvars, invars = _open_jaxpr_with_split(2)
+        assert split_jaxpr_invars(jaxpr, len(invars)) == (constvars, invars)
+        assert jaxpr_constvars(jaxpr, len(invars)) == constvars
+
+    def test_the_recovered_consts_actually_evaluate(self):
+        # The regression that broke the whole suite on jaxlib 0.11.1: reading
+        # ``jaxpr.constvars`` back handed eval_jaxpr too few values.
+        jaxpr, _, invars = _open_jaxpr_with_split(2)
+        consts = [jnp.full((3,), 2.0), jnp.full((3,), 5.0)]
+        out = jax.core.eval_jaxpr(jaxpr, consts, jnp.full((3,), 3.0))
+        assert jnp.allclose(out[0], 2.0 * 3.0 + 5.0)
+
+    def test_zero_constvars_edge(self):
+        jaxpr, constvars, invars = _open_jaxpr_with_split(0)
+        assert constvars == []
+        assert jaxpr_constvars(jaxpr, len(invars)) == []
+
+    def test_all_constvars_edge(self):
+        jaxpr, constvars, invars = _open_jaxpr_with_split(3)
+        assert invars == []
+        assert jaxpr_constvars(jaxpr, 0) == constvars
+
+    @pytest.mark.parametrize('num_invars', [-1, 4])
+    def test_out_of_range_num_invars_raises(self, num_invars):
+        # A miscounted call must fail at the compiler boundary, not silently
+        # produce a misaligned eval_jaxpr argument list.
+        jaxpr, _, _ = _open_jaxpr_with_split(1)
+        with pytest.raises(ValueError, match='positional inputs'):
+            split_jaxpr_invars(jaxpr, num_invars)
+
+
 class TestPrimitive:
     def test_jit(self):
         @jit

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/graph.py

@@ -22,6 +22,7 @@ import brainstate
 import jax
 import numpy as np
 
+from braintrace._compatible_imports import jaxpr_all_invars
 from braintrace._misc import NotSupportedError
 from braintrace._typing import (
     ETraceVals,
@@ -571,7 +572,7 @@ def compile_etrace_graph(
             if relation.control_flow_context is not None:
                 continue
             for hpo_jaxpr in relation.y_to_hidden_group_jaxprs:
-                for v in hpo_jaxpr.invars + hpo_jaxpr.constvars:
+                for v in jaxpr_all_invars(hpo_jaxpr):
                     out_wy2hid_jaxvars_dict[v] = None
         out_wy2hid_jaxvars = list(out_wy2hid_jaxvars_dict)
 

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/hid_param_op.py

@@ -63,6 +63,7 @@ from braintrace._compatible_imports import (
     is_scan_primitive,
     is_while_primitive,
     is_cond_primitive,
+    jaxpr_constvars,
 )
 from braintrace._op import (
     get_trainable_invars,
@@ -248,7 +249,9 @@ class HiddenParamOpRelation(NamedTuple):
         """
         vals_of_hidden_groups = []
         for jaxpr, group in zip(self.y_to_hidden_group_jaxprs, self.hidden_groups):
-            consts = [const_vals[var] for var in jaxpr.constvars]
+            # ``y`` is the jaxpr's only invar (see ``_build_transition_jaxpr``);
+            # everything ahead of it is a constvar bound from the forward pass.
+            consts = [const_vals[var] for var in jaxpr_constvars(jaxpr, 1)]
             # ``list[Array]`` before concatenation, a single ``Array`` after.
             hidden_vals: Any = jax.core.eval_jaxpr(jaxpr, consts, y_val)
             if concat_hidden_vals:

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/hid_param_op_test.py

@@ -195,3 +195,32 @@ class TestTrainableInvarsPopulatedForDense:
         assert r.trainable_leaf_indices['weight'] == 0
         assert r.trainable_processing_chains['weight'] == ()
         assert r.trainable_param_states['weight'] is not None
+
+
+class TestYToHiddenJaxprConstvarSplit:
+    """The ``y -> hidden group`` jaxprs keep a recoverable const/invar split.
+
+    Since JAX 0.11 an open jaxpr with symbolic-only constvars reports
+    ``constvars == []``, so the split is recovered from the invar count instead:
+    these jaxprs take exactly one invar, the relation's ``y_var``. The
+    executable counterpart of this assertion lives in
+    ``_testing/scenario_catalog_test.py::test_gru_y_to_hidden_groups_executes``.
+    """
+
+    def test_y_is_the_sole_invar_and_the_rest_are_constvars(self):
+        import jax.numpy as jnp
+        from braintrace._compatible_imports import jaxpr_all_invars, split_jaxpr_invars
+
+        gru = braintrace.nn.GRUCell(3, 4)
+        brainstate.nn.init_all_states(gru)
+        relations = find_hidden_param_op_relations_from_module(gru, jnp.zeros(3))
+
+        assert len(relations) > 0
+        checked = 0
+        for r in relations:
+            for jaxpr in r.y_to_hidden_group_jaxprs:
+                constvars, invars = split_jaxpr_invars(jaxpr, 1)
+                assert invars == [r.y_var]
+                assert constvars == jaxpr_all_invars(jaxpr)[:-1]
+                checked += 1
+        assert checked > 0

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/hidden_group.py

@@ -63,6 +63,7 @@ from braintrace._compatible_imports import (
     is_scan_primitive,
     is_while_primitive,
     is_cond_primitive,
+    jaxpr_constvars,
     scan_num_consts_carry,
 )
 from braintrace._op import is_etp_primitive, is_etp_enable_gradient_primitive
@@ -1621,6 +1622,11 @@ class JaxprEvalForHiddenGroup(JaxprEvaluation):
         pattern = build_snap_pattern(
             group.transition_jaxpr, group.varshape, self.sparse_n,
             num_state=group.num_state,
+            # Seed the adjacency walk from the hidden invars explicitly. Since
+            # JAX 0.11 the jaxpr's own ``invars`` also lists the constvars, so
+            # the implicit default would widen the seed set and inflate the
+            # derived pattern (usually all the way to conservative).
+            hidden_invars=group.hidden_invars,
             max_jacobian_elements=self.snap_max_jacobian_elements,
         )
         if pattern.conservative:
@@ -1764,7 +1770,10 @@ class JaxprEvalForHiddenGroup(JaxprEvaluation):
                     for outvar in hidden_outvars
                 ],
                 transition_jaxpr=jaxpr,
-                transition_jaxpr_constvars=list(jaxpr.constvars),
+                # ``write_jaxpr_of_hidden_group_transition`` gives the jaxpr one
+                # invar per hidden state; the leading remainder are the
+                # constvars bound from the forward pass.
+                transition_jaxpr_constvars=jaxpr_constvars(jaxpr, len(hidden_invars)),
                 is_diagonal_recurrence=not self.include_recurrent_mixing,
             )
             group = self._attach_snap_pattern(group)
@@ -1832,7 +1841,10 @@ class JaxprEvalForHiddenGroup(JaxprEvaluation):
                 hidden_paths=[self.outvar_to_hidden_path[outvar]],
                 hidden_states=[self.path_to_state[self.outvar_to_hidden_path[outvar]]],
                 transition_jaxpr=zero_jaxpr,
-                transition_jaxpr_constvars=list(zero_jaxpr.constvars),
+                # Mirrors ``zero_jaxpr``'s ``constvars=[outvar], invars=[invar]``
+                # above -- read literally, not back off the jaxpr, which since
+                # JAX 0.11 no longer records the boundary for a value-less split.
+                transition_jaxpr_constvars=[outvar],
                 # A zero-recurrence transition (``D^t = 0``) is trivially diagonal;
                 # keep the flag mode-derived for uniformity (this fallback only
                 # fires in the default mode in practice).

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/hidden_group_test.py

@@ -1975,3 +1975,46 @@ class TestSplitHiddenWidthGuard:
         assert len(parts) == 2
         np.testing.assert_array_equal(np.asarray(parts[0]), np.ones((2, 3)))
         np.testing.assert_array_equal(np.asarray(parts[1]), np.full((2, 3), 2.0))
+
+
+from braintrace._compatible_imports import jaxpr_all_invars
+
+
+class TestTransitionJaxprConstvarSplit:
+    """A compiled group's const/invar split stays intact on every JAX version.
+
+    JAX 0.11 merged ``ClosedJaxpr`` into ``Jaxpr`` and stopped recording the
+    boundary for an open jaxpr with symbolic-only constvars, which is exactly
+    what the compiler builds here. ``transition_jaxpr_constvars`` must therefore
+    be carried explicitly rather than read back off the jaxpr.
+    """
+
+    def _gru_group(self):
+        gru = braintrace.nn.GRUCell(3, 4)
+        brainstate.nn.init_all_states(gru)
+        groups, _ = find_hidden_groups_from_module(gru, brainstate.random.randn(3))
+        assert len(groups) == 1
+        return groups[0]
+
+    def test_constvars_are_the_leading_positional_inputs(self):
+        group = self._gru_group()
+        all_in = jaxpr_all_invars(group.transition_jaxpr)
+        n_const = len(all_in) - len(group.hidden_invars)
+        # A non-trivial split: reading ``.constvars`` back would give [] on 0.11.
+        assert n_const > 0
+        assert group.transition_jaxpr_constvars == all_in[:n_const]
+        assert list(group.hidden_invars) == all_in[n_const:]
+
+    def test_the_transition_evaluates_with_those_constvars(self):
+        # The arity failure that broke the suite on jaxlib 0.11.1 surfaces here.
+        group = self._gru_group()
+        input_vals = [
+            jnp.zeros(v.aval.shape, v.aval.dtype)
+            for v in group.transition_jaxpr_constvars
+        ]
+        hidden_vals = [
+            jnp.zeros(v.aval.shape, v.aval.dtype)
+            for v in group.hidden_invars
+        ]
+        out = group.transition(hidden_vals, input_vals)
+        assert len(out) == len(group.hidden_outvars)

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/scan_descent.py

@@ -64,6 +64,7 @@ from braintrace._compatible_imports import (
     is_cond_primitive,
     is_scan_primitive,
     is_while_primitive,
+    jaxpr_all_invars,
     new_var,
     scan_num_consts_carry,
     scan_params_add_ys,
@@ -443,7 +444,7 @@ def analyze_and_rewrite_scan(eqn: JaxprEqn, minfo: 'ModuleInfo') -> Optional[Sca
         if r.x_var is not None:
             hoist_seen[r.x_var] = None
         for j in r.y_to_hidden_group_jaxprs:
-            for v in list(j.invars) + list(j.constvars):
+            for v in jaxpr_all_invars(j):
                 hoist_seen[v] = None
     for g in body_groups:
         for v in list(g.hidden_invars) + list(g.transition_jaxpr_constvars):

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/scan_descent_test.py

@@ -13,7 +13,12 @@ import pytest
 
 import braintrace
 from braintrace import ControlFlowPolicy
-from braintrace._compatible_imports import is_scan_primitive, scan_num_consts_carry
+from braintrace._compatible_imports import (
+    is_scan_primitive,
+    jaxpr_constvars,
+    scan_num_consts_carry,
+    split_jaxpr_invars,
+)
 from braintrace._compiler.module_info import extract_module_info
 from braintrace._compiler.scan_descent import (
     _descent_blockers,
@@ -219,8 +224,13 @@ class TestAnalyzeAndRewriteScan:
         # outer-facing hidden vars are the scan carry vars, known to minfo
         assert g.hidden_outvars[0] in minfo.outvar_to_hidden_path
         assert g.hidden_invars[0] in minfo.invar_to_hidden_path
-        # body-scoped transition: one substep, invars == descent.body_hidden_invars
-        assert list(g.transition_jaxpr.invars) == list(g.descent.body_hidden_invars)
+        # body-scoped transition: one substep, invars == descent.body_hidden_invars.
+        # Split by invar count rather than reading ``.invars`` -- since JAX 0.11
+        # that also lists the constvars for a value-less open jaxpr.
+        _, transition_invars = split_jaxpr_invars(
+            g.transition_jaxpr, len(g.descent.body_hidden_invars)
+        )
+        assert list(transition_invars) == list(g.descent.body_hidden_invars)
 
         assert len(bundle.relations) == 1
         r = bundle.relations[0]
@@ -231,7 +241,7 @@ class TestAnalyzeAndRewriteScan:
         m = bundle.info.stacked_var_map
         assert r.x_var in m and r.y_var in m
         for j in r.y_to_hidden_group_jaxprs:
-            assert all(v in m for v in j.constvars)
+            assert all(v in m for v in jaxpr_constvars(j, 1))
         assert all(v in m for v in g.descent.body_hidden_invars)
         assert all(v in m for v in g.transition_jaxpr_constvars)
         # stacked avals carry the substep axis

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_compiler/tests/compiler_oracle_test.py

@@ -47,6 +47,7 @@ import jax.numpy as jnp
 import numpy as np
 
 from braintrace import compile_etrace_graph
+from braintrace._compatible_imports import jaxpr_constvars
 from braintrace._testing.scenario_catalog import (
     ElemwiseOnlyRNN,
     PartialPathRNN,
@@ -67,8 +68,11 @@ def _transition_callable(rel, group, const_vals):
     The returned callable takes ``y_val`` and returns the *concatenated*
     hidden state value for the matching group.
     """
-    consts = [const_vals[v] for v in rel.y_to_hidden_group_jaxprs[0].constvars]
     jaxpr = rel.y_to_hidden_group_jaxprs[0]
+    # ``y`` is the jaxpr's sole invar; everything ahead of it is a constvar.
+    # Read via the helper, not ``jaxpr.constvars`` -- since JAX 0.11 an open
+    # jaxpr with symbolic-only constvars reports that as empty.
+    consts = [const_vals[v] for v in jaxpr_constvars(jaxpr, 1)]
 
     def f(y_val):
         out = jax.core.eval_jaxpr(jaxpr, consts, y_val)

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace/_version.py

@@ -13,5 +13,5 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.2.5"
+__version__ = "0.2.6"
 __version_info__ = tuple(map(int, __version__.split('.')))

{braintrace-0.2.5 → braintrace-0.2.6}/braintrace.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.5
+Version: 0.2.6
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license