diff --git a/braintrace/_algorithm/base.py b/braintrace/_algorithm/base.py index b403164e..7032d449 100644 --- a/braintrace/_algorithm/base.py +++ b/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""" diff --git a/braintrace/_algorithm/io_dim_vjp.py b/braintrace/_algorithm/io_dim_vjp.py index deb9d640..b83b225d 100644 --- a/braintrace/_algorithm/io_dim_vjp.py +++ b/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 diff --git a/braintrace/_algorithm/io_dim_vjp_test.py b/braintrace/_algorithm/io_dim_vjp_test.py index 23640097..1ce0e75d 100644 --- a/braintrace/_algorithm/io_dim_vjp_test.py +++ b/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) diff --git a/braintrace/_compatible_imports.py b/braintrace/_compatible_imports.py index 690ec26a..4e5d32cf 100644 --- a/braintrace/_compatible_imports.py +++ b/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' diff --git a/braintrace/_compiler/base.py b/braintrace/_compiler/base.py index 5eeb2753..fef94320 100644 --- a/braintrace/_compiler/base.py +++ b/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,7 +39,7 @@ def find_matched_vars( - invars: Sequence[Var], + invars: Sequence[Union[Var, Literal]], invar_needed_in_oth_eqns: Container[Var] ) -> List[Var]: """ @@ -46,8 +47,8 @@ def find_matched_vars( 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,7 +66,7 @@ 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: """ @@ -73,8 +74,8 @@ def find_element_exist_in_the_set( 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. diff --git a/braintrace/_compiler/canonicalize.py b/braintrace/_compiler/canonicalize.py index 9be96a8e..bcccb931 100644 --- a/braintrace/_compiler/canonicalize.py +++ b/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]: diff --git a/braintrace/_compiler/hid_param_op.py b/braintrace/_compiler/hid_param_op.py index 5cd0d65f..2b9fd59b 100644 --- a/braintrace/_compiler/hid_param_op.py +++ b/braintrace/_compiler/hid_param_op.py @@ -49,6 +49,7 @@ Sequence, Set, Tuple, + Union, ) import brainstate @@ -57,6 +58,7 @@ from braintrace._compatible_imports import ( Primitive, Var, + Literal, JaxprEqn, Jaxpr, is_jit_primitive, @@ -64,6 +66,7 @@ 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) diff --git a/braintrace/_compiler/hidden_pertubation.py b/braintrace/_compiler/hidden_pertubation.py index 52c5ade2..898fcc8a 100644 --- a/braintrace/_compiler/hidden_pertubation.py +++ b/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,6 +23,7 @@ from braintrace._compatible_imports import ( Var, + Literal, JaxprEqn, Jaxpr, ClosedJaxpr, @@ -30,6 +31,7 @@ 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]: diff --git a/braintrace/_compiler/module_info.py b/braintrace/_compiler/module_info.py index 3732d3e8..1463b6cc 100644 --- a/braintrace/_compiler/module_info.py +++ b/braintrace/_compiler/module_info.py @@ -468,8 +468,10 @@ def _process( # - Outputs # 2. "jaxpr_outs[self.num_out:]" corresponds to extra output in "augmented_jaxpr" # - others + # (typed by variable: consumers look temporaries up by ``Var``; a + # literal state outvar is never looked up) temps = { - v: r for v, r in + cast(Var, v): r for v, r in zip( self.jaxpr.outvars[self.num_var_out:], jaxpr_outs[self.num_var_out:] diff --git a/braintrace/_compiler/output_bypass.py b/braintrace/_compiler/output_bypass.py index 32030757..fbdf4614 100644 --- a/braintrace/_compiler/output_bypass.py +++ b/braintrace/_compiler/output_bypass.py @@ -56,7 +56,7 @@ the effect). """ -from typing import Any, Dict, FrozenSet, Hashable, List, NamedTuple, Sequence, Set, Tuple, cast +from typing import Any, Dict, FrozenSet, Hashable, List, NamedTuple, Sequence, Set, Tuple import jax import jax.numpy as jnp @@ -69,6 +69,7 @@ is_while_primitive, new_jaxpr_eqn, new_var, + shaped_aval, stop_gradient_p, ) from .hidden_group import _transition_contains_while @@ -117,8 +118,8 @@ def _is_blocking(eqn: JaxprEqn) -> bool: def _shaped(v: Var) -> jax.core.ShapedArray: - # every variable of a traced step is array-valued - aval = cast(jax.core.ShapedArray, v.aval) + # a fresh aval of the same shape and dtype (no weak type) + aval = shaped_aval(v) return jax.core.ShapedArray(aval.shape, aval.dtype) diff --git a/braintrace/_compiler/scan_descent.py b/braintrace/_compiler/scan_descent.py index 670d8f88..bd37d991 100644 --- a/braintrace/_compiler/scan_descent.py +++ b/braintrace/_compiler/scan_descent.py @@ -51,7 +51,7 @@ ETP mixing into a body transition has no consumer. """ -from typing import Dict, List, NamedTuple, Optional, Sequence, Set, Tuple, TYPE_CHECKING +from typing import Dict, List, NamedTuple, Optional, Sequence, Set, Tuple, TYPE_CHECKING, cast from braintrace._compatible_imports import ( ClosedJaxpr, @@ -65,6 +65,7 @@ new_var, scan_num_consts_carry, scan_params_add_ys, + shaped_aval, ) from braintrace._op import is_etp_primitive from braintrace._typing import Path @@ -266,7 +267,7 @@ def add_scan_ys( ) new_closed = ClosedJaxpr(new_body, body_closed.consts) stacked = { - v: new_var('', v.aval.update(shape=(length, *v.aval.shape))) + v: new_var('', shaped_aval(v).update(shape=(length, *shaped_aval(v).shape))) for v in extra } new_params = scan_params_add_ys( @@ -371,11 +372,13 @@ def analyze_and_rewrite_scan(eqn: JaxprEqn, minfo: 'ModuleInfo') -> Optional[Sca body_invar_to_hidden_path = { body.invars[num_consts + c]: p for c, p in carry_hidden.items() } + # (typed as variables: a hidden carry output is used as a variable key) + body_outvars = cast(List[Var], list(body.outvars)) body_outvar_to_hidden_path = { - body.outvars[c]: p for c, p in carry_hidden.items() + body_outvars[c]: p for c, p in carry_hidden.items() } body_hidden_outvar_to_invar = { - body.outvars[c]: body.invars[num_consts + c] + body_outvars[c]: body.invars[num_consts + c] for c in carry_hidden } body_invar_to_weight_path: Dict[Var, Path] = {} @@ -465,7 +468,8 @@ def analyze_and_rewrite_scan(eqn: JaxprEqn, minfo: 'ModuleInfo') -> Optional[Sca path_to_carry = {p: c for c, p in carry_hidden.items()} final_groups = [] for g in body_groups: - outer_invars = [eqn.invars[num_consts + path_to_carry[p]] + # the v1 guard above checked every hidden carry init is a ``Var`` + outer_invars = [cast(Var, eqn.invars[num_consts + path_to_carry[p]]) for p in g.hidden_paths] outer_outvars = [eqn.outvars[path_to_carry[p]] for p in g.hidden_paths] final_groups.append(g._replace( diff --git a/docs/specs/2026-10-24-mypy-typecheck.md b/docs/specs/2026-10-24-mypy-typecheck.md new file mode 100644 index 00000000..601d8d52 --- /dev/null +++ b/docs/specs/2026-10-24-mypy-typecheck.md @@ -0,0 +1,49 @@ +# Fix the CI type check (mypy) + +## Problem + +The `typecheck_and_build` CI job fails on `main`: `python -m mypy braintrace` +reports 36 errors in 10 files. Almost all of them come from the JAX stubs, +which now type: + +- `Var.aval` as the abstract base `AbstractValue`, which has no `shape`, + `dtype` or `ndim`; +- `JaxprEqn.invars` / `Jaxpr.outvars` as `list[Var | Literal]`; +- `Var(...)` with a single argument (the pre-0.6.2 two-argument form is a + `call-arg` error). + +The other error is `FlattedDict.split`, whose stub returns a `NestedDict` +while it returns one `FlattedDict` per filter at runtime. + +## Fix + +Type-only, apart from the latent bug below. + +- `_compatible_imports.shaped_aval(atom)`: a `cast` to + `jax.core.ShapedArray`. Every atom of a traced array program carries one. + It is used at every `.aval.shape` / `.dtype` / `.ndim` site mypy flags. +- The membership helpers `find_matched_vars` and + `find_element_exist_in_the_set` accept `Var | Literal`. They already + filtered with `isinstance(..., Var)`. +- `hid_param_op`: trainable invars are `Var | Literal`. A literal never + traces to a `ParamState` (its path is `None`), and the loop narrows on + that. +- Explicit annotations or justified casts for the remaining sites: the + detach list in `_detach_invars`, the scan descent body maps (hidden carry + inits are guarded to be `Var`s), the module temporaries, and the + `FlattedDict.split` result. +- `new_var`: `# type: ignore[call-arg]` on the old-JAX branch only. + +## Latent bug fixed on the way + +`IODimVjpAlgorithm.get_etrace_of` tested `etrace_x_key(relation.x_var) is not +None`. `etrace_x_key` returns `id(x_var)` and is never `None`, so for an +element-wise relation (`x_var is None`) it looked up `id(None)` and raised +`KeyError`. It now checks `relation.x_var is not None`. + +## Tests + +- `python -m mypy braintrace`: `Success: no issues found`. +- `io_dim_vjp_test.py::TestStateLifecycle::test_get_etrace_of_elementwise_weight` + fails with `KeyError` before the fix and passes after. +- The full suite passes unchanged.