Fix the CI type check (mypy) and get_etrace_of on element-wise weights

braintrace/_algorithm/base.py

@@ -21,7 +21,7 @@
 
 from __future__ import annotations
 
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, Tuple, cast
 
 import brainstate
 
@@ -257,11 +257,16 @@ def _split_state(self) -> None:
         #   - the weight states, which is invariant during the training process
         #   - the hidden states, the recurrent states, which may be changed between different training epochs
         #   - the other states, which may be changed between different training epochs
+        # ``FlattedDict.split`` returns one ``FlattedDict`` per filter (the
+        # stubs only promise a ``NestedDict``)
         (
             self._param_states,
             self._hidden_states,
             self._other_states
-        ) = self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...)
+        ) = cast(
+            Tuple[brainstate.util.FlattedDict, brainstate.util.FlattedDict, brainstate.util.FlattedDict],
+            self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...),
+        )
 
     def compile_graph(self, *args: Any) -> None:
         r"""

braintrace/_algorithm/io_dim_vjp.py

@@ -34,6 +34,7 @@
 import jax.numpy as jnp
 import brainunit as u
 
+from braintrace._compatible_imports import shaped_aval
 from braintrace._compiler import ControlFlowPolicy, HiddenGroup, HiddenParamOpRelation
 from braintrace._op import (
     etp_elemwise_p,
@@ -262,8 +263,8 @@ def _init_IO_dim_state(
         if x_key not in etrace_xs:
             x_repr_fn = get_pp_x_repr(relation.primitive)
             if x_repr_fn is None:
-                shape = relation.x_var.aval.shape
-                dtype = relation.x_var.aval.dtype
+                shape = shaped_aval(relation.x_var).shape
+                dtype = shaped_aval(relation.x_var).dtype
             else:
                 # the trace filters the primitive's x *representation*
                 # (e.g. embedding: the one-hot encoding of its integer
@@ -272,12 +273,12 @@ def _init_IO_dim_state(
                 x_aval = jax.eval_shape(
                     lambda x_, _fn=x_repr_fn, _w=weight_avals: _fn(x_, _w),
                     jax.ShapeDtypeStruct(
-                        relation.x_var.aval.shape, relation.x_var.aval.dtype),
+                        shaped_aval(relation.x_var).shape, shaped_aval(relation.x_var).dtype),
                 )
                 shape, dtype = x_aval.shape, x_aval.dtype
             etrace_xs[x_key] = EligibilityTrace(u.math.zeros(shape, dtype))
 
-    y_shape = relation.y_var.aval.shape
+    y_shape = shaped_aval(relation.y_var).shape
     group: HiddenGroup
     for group in relation.hidden_groups:
         # Exact match required, or (elemwise only) allow trailing-dim match
@@ -613,7 +614,7 @@ def _call(df_: Any, w_: Any, _rule: Any = xy_to_dw_rule, _params: Any = eqn_para
             # gate the branch.)
             elemwise_batched = (
                 relation.primitive is etp_elemwise_p
-                and len(group.varshape) > relation.y_var.aval.ndim
+                and len(group.varshape) > shaped_aval(relation.y_var).ndim
             )
 
             if fast_solve:
@@ -926,8 +927,9 @@ def get_etrace_of(self, weight: brainstate.ParamState | Path) -> Tuple[Dict, Dic
             find_this_weight = True
 
             # get the weight_op input
-            wx_var = etrace_x_key(relation.x_var)
-            if wx_var is not None:
+            # (an element-wise relation has no x trace)
+            if relation.x_var is not None:
+                wx_var = etrace_x_key(relation.x_var)
                 etrace_xs[wx_var] = self.etrace_xs[wx_var].value
 
             # get the weight_op df

braintrace/_algorithm/io_dim_vjp_test.py

@@ -247,6 +247,27 @@ def test_get_etrace_of_returns_xs_dfs_tuple(self):
         assert isinstance(result, tuple)
         assert len(result) == 2
 
+    def test_get_etrace_of_elementwise_weight(self):
+        # An element-wise relation has no x trace; only its df traces are returned.
+        class Leaky(brainstate.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.w = brainstate.ParamState(jnp.full((3,), 0.5))
+
+            def init_state(self, *args, **kwargs):
+                self.h = brainstate.HiddenState(jnp.zeros(3))
+
+            def update(self, x):
+                self.h.value = 0.9 * self.h.value + braintrace.element_wise(self.w.value) * x
+                return self.h.value
+
+        model = Leaky()
+        brainstate.nn.init_all_states(model)
+        algo = _compiled(model)
+        xs, dfs = algo.get_etrace_of(model.w)
+        assert xs == {}
+        assert len(dfs) == 1
+
     def test_get_etrace_of_plain_weight_raises(self):
         model = _build(om.tanh_rnn)
         algo = _compiled(model)

braintrace/_compatible_imports.py

@@ -23,6 +23,7 @@
     'ClosedJaxpr',
     'Literal',
     'new_var',
+    'shaped_aval',
     'new_jaxpr_eqn',
     'stop_gradient_p',
     'is_jit_primitive',
@@ -37,7 +38,7 @@
     'jaxpr_constvars',
 ]
 
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple, cast
 
 from brainstate._compatible_import import Primitive, Var, JaxprEqn, Jaxpr, ClosedJaxpr, Literal, wrap_init
 
@@ -67,11 +68,32 @@
 
 def new_var(suffix: Any, aval: Any) -> Var:
     if jax.__version_info__ < (0, 6, 2):
-        return Var(suffix, aval)
+        # the pre-0.6.2 constructor; the installed stubs only know ``Var(aval)``
+        return Var(suffix, aval)  # type: ignore[call-arg]
     else:
         return Var(aval)
 
 
+def shaped_aval(atom: Any) -> jax.core.ShapedArray:
+    """Return the aval of a jaxpr variable or literal as a ``ShapedArray``.
+
+    The JAX stubs type ``Var.aval`` as the abstract base ``AbstractValue``,
+    which has no ``shape`` / ``dtype`` / ``ndim``. Every atom of a traced
+    array program carries a ``ShapedArray``, so this only narrows the type.
+
+    Parameters
+    ----------
+    atom : Var or Literal
+        A jaxpr atom.
+
+    Returns
+    -------
+    jax.core.ShapedArray
+        Its aval.
+    """
+    return cast(jax.core.ShapedArray, atom.aval)
+
+
 def is_jit_primitive(eqn: JaxprEqn) -> bool:
     if jax.__version_info__ < (0, 7, 0):
         return eqn.primitive.name == 'pjit'

braintrace/_compiler/base.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Container, Dict, Sequence, Set, List
+from typing import Container, Dict, Sequence, Set, List, Union
 
 from braintrace._compatible_imports import (
     Var,
+    Literal,
     Jaxpr,
     JaxprEqn,
     is_jit_primitive,
@@ -38,16 +39,16 @@
 
 
 def find_matched_vars(
-    invars: Sequence[Var],
+    invars: Sequence[Union[Var, Literal]],
     invar_needed_in_oth_eqns: Container[Var]
 ) -> List[Var]:
     """
     Checking whether the invars are matched with the invar_needed_in_oth_eqns.
 
     Parameters
     ----------
-    invars : Sequence[Var]
-        The input variables of the equation.
+    invars : Sequence[Var or Literal]
+        The input atoms of the equation; literals never match.
     invar_needed_in_oth_eqns : Container[Var]
         The variables needed in the other equations (a set, or an
         insertion-ordered dict used as an ordered set).
@@ -65,16 +66,16 @@ def find_matched_vars(
 
 
 def find_element_exist_in_the_set(
-    elements: Sequence[Var],
+    elements: Sequence[Union[Var, Literal]],
     the_set: Set[Var]
 ) -> Var | None:
     """
     Checking whether the jaxpr vars contain the weight variables.
 
     Parameters
     ----------
-    elements : Sequence[Var]
-        The input variables of the equation.
+    elements : Sequence[Var or Literal]
+        The input atoms of the equation; literals never match.
     the_set : Set[Var]
         The set of the weight variables.
 

braintrace/_compiler/canonicalize.py

@@ -64,6 +64,7 @@
     new_jaxpr_eqn,
     new_var,
     scan_num_consts_carry,
+    shaped_aval,
 )
 from braintrace._misc import CompilationError
 from braintrace._op import is_etp_primitive
@@ -985,7 +986,7 @@ def resolve(atom: Any) -> Any:
             outvar_map = final_carry_map if t == final_t else {}
             for body_eqn in body.eqns:
                 fresh_outvars = [
-                    outvar_map.get(v) if isinstance(v, Var) and v in outvar_map
+                    outvar_map[v] if isinstance(v, Var) and v in outvar_map
                     else fresh_like(v)
                     for v in body_eqn.outvars
                 ]
@@ -1007,7 +1008,7 @@ def resolve(atom: Any) -> Any:
         for ov, final_atom in zip(eqn.outvars[:num_carry], carry_atoms):
             if ov not in consumed or final_atom is ov:
                 continue
-            shape = tuple(ov.aval.shape)
+            shape = tuple(shaped_aval(ov).shape)
             new_eqns.append(new_jaxpr_eqn(
                 [final_atom],
                 [ov],
@@ -1023,7 +1024,7 @@ def resolve(atom: Any) -> Any:
         for j, ov in enumerate(eqn.outvars[num_carry:]):
             if ov not in consumed:
                 continue
-            stacked_shape = tuple(ov.aval.shape)
+            stacked_shape = tuple(shaped_aval(ov).shape)
             slice_shape = stacked_shape[1:]
             slice_vars = []
             for y_atom in ys_atoms[j]:

braintrace/_compiler/hid_param_op.py

@@ -49,6 +49,7 @@
     Sequence,
     Set,
     Tuple,
+    Union,
 )
 
 import brainstate
@@ -57,13 +58,15 @@
 from braintrace._compatible_imports import (
     Primitive,
     Var,
+    Literal,
     JaxprEqn,
     Jaxpr,
     is_jit_primitive,
     is_scan_primitive,
     is_while_primitive,
     is_cond_primitive,
     jaxpr_constvars,
+    shaped_aval,
 )
 from braintrace._op import (
     get_trainable_invars,
@@ -287,8 +290,8 @@ def __repr__(self) -> str:
 
 def _resolve_eqn_trainable_invars(
     eqn: JaxprEqn,
-) -> Dict[str, Var]:
-    """Return ``{key: invar_var}`` for every trainable input of *eqn*."""
+) -> Dict[str, Union[Var, Literal]]:
+    """Return ``{key: invar_atom}`` for every trainable input of *eqn*."""
     key_to_idx = get_trainable_invars(eqn.primitive, eqn.params)
     return {k: eqn.invars[i] for k, i in key_to_idx.items()}
 
@@ -327,10 +330,10 @@ def _resolve_eqn_vars(
 # ---------------------------------------------------------------------------
 
 def _trace_var_to_param(
-    var: Var,
+    var: Union[Var, Literal],
     producers: Dict[Var, JaxprEqn],
     invar_to_weight_path: Dict[Var, Path],
-    cache: Optional[Dict[Var, Optional[Tuple[Path, Tuple[Primitive, ...]]]]] = None,
+    cache: Optional[Dict[Any, Optional[Tuple[Path, Tuple[Primitive, ...]]]]] = None,
 ) -> Tuple[Optional[Path], Tuple[Primitive, ...]]:
     """Trace *var* backward through the Jaxpr to its originating ``ParamState``.
 
@@ -883,7 +886,8 @@ def find_hidden_param_op_relations_from_jaxpr(
                 invar, producers, invar_to_weight_path,
                 cache=weight_trace_cache,
             )
-            if t_path is None:
+            # a literal never traces to a ParamState (its path is None)
+            if t_path is None or not isinstance(invar, Var):
                 # Trainable invar doesn't trace to any ParamState (e.g. a
                 # constant bias passed directly as a jnp.array). Emit an INFO
                 # diagnostic so users know no gradient will be produced for
@@ -945,23 +949,23 @@ def find_hidden_param_op_relations_from_jaxpr(
         path_class: Dict[Path, str] = {}
         for hvar in list(reachable_hvars):
             try:
-                jax.numpy.broadcast_shapes(y_var.aval.shape, hvar.aval.shape)
+                jax.numpy.broadcast_shapes(shaped_aval(y_var).shape, shaped_aval(hvar).shape)
             except ValueError:
                 emit(
                     kind=DiagnosticKind.RELATION_EXCLUDED_SHAPE_MISMATCH,
                     level=DiagnosticLevel.WARNING,
                     message=(
                         f'ETP op {primitive.name}: weight={weight_path}, '
-                        f'y shape={y_var.aval.shape} not broadcastable with '
-                        f'hidden shape={hvar.aval.shape} at '
+                        f'y shape={shaped_aval(y_var).shape} not broadcastable with '
+                        f'hidden shape={shaped_aval(hvar).shape} at '
                         f'{outvar_to_hidden_path[hvar]}. Removing connection.'
                     ),
                     primitive=primitive,
                     weight_path=weight_path,
                     hidden_paths=(outvar_to_hidden_path[hvar],),
                     context={
-                        'y_shape': tuple(y_var.aval.shape),
-                        'hidden_shape': tuple(hvar.aval.shape),
+                        'y_shape': tuple(shaped_aval(y_var).shape),
+                        'hidden_shape': tuple(shaped_aval(hvar).shape),
                     },
                 )
                 reachable_hvars.pop(hvar, None)

braintrace/_compiler/hidden_pertubation.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 
-from typing import Dict, FrozenSet, List, Mapping, Optional, Set, Sequence, NamedTuple, Any, cast
+from typing import Dict, FrozenSet, List, Mapping, Optional, Set, Sequence, NamedTuple, Any, Union
 
 import brainstate
 import jax.core
@@ -23,13 +23,15 @@
 
 from braintrace._compatible_imports import (
     Var,
+    Literal,
     JaxprEqn,
     Jaxpr,
     ClosedJaxpr,
     new_var,
     new_jaxpr_eqn,
     stop_gradient_p,
     is_while_primitive,
+    shaped_aval,
 )
 from braintrace._misc import (
     git_issue_addr,
@@ -185,7 +187,7 @@ def init_perturb_data(self) -> Sequence[jax.Array]:
             One zero array per perturbation variable, matching its shape and
             dtype.
         """
-        return [jax.numpy.zeros(v.aval.shape, dtype=v.aval.dtype) for v in self.perturb_vars]
+        return [jax.numpy.zeros(shaped_aval(v).shape, dtype=shaped_aval(v).dtype) for v in self.perturb_vars]
 
     def bypass_hidden_data_to_hidden_group_data(
         self,
@@ -225,7 +227,7 @@ def bypass_hidden_data_to_hidden_group_data(
         by_path = dict(zip(self.bypass_hidden_paths, bypass_hidden_data))
         per_perturb = []
         for path, v in zip(self.perturb_hidden_paths, self.perturb_vars):
-            aval = cast(jax.core.ShapedArray, v.aval)
+            aval = shaped_aval(v)
             per_perturb.append(by_path[path] if path in by_path else jnp.zeros(aval.shape, aval.dtype))
         return self.perturb_data_to_hidden_group_data(per_perturb, hidden_groups)
 
@@ -610,7 +612,7 @@ def _detach_invars(self, eqn: JaxprEqn) -> JaxprEqn:
         returns a copy of *eqn* consuming the detached variables. Literal
         inputs are left untouched.
         """
-        new_invars = []
+        new_invars: List[Union[Var, Literal]] = []
         n_detached = 0
         for iv in eqn.invars:
             if not isinstance(iv, Var):
@@ -693,7 +695,7 @@ def _eval_eqn(self, eqn: JaxprEqn) -> None:
             )
 
     def _new_var_like(self, v: Var) -> Var:
-        return new_var('', jax.core.ShapedArray(v.aval.shape, v.aval.dtype))
+        return new_var('', jax.core.ShapedArray(shaped_aval(v).shape, shaped_aval(v).dtype))
 
 
 def _while_eqns_reaching(jaxpr: Jaxpr, targets: Set[Var]) -> Set[int]: