[chaoming0625]

Fixes the failing typecheck_and_build CI job. On main, python -m mypy braintrace reported 36 errors in 10 files. It is now Success: no issues found in 63 source files.

Cause

The newer JAX stubs type:

• Var.aval as the abstract base AbstractValue (no shape / dtype / ndim);
• eqn.invars / jaxpr.outvars as list[Var | Literal];
• Var(...) with one argument.

Also, the FlattedDict.split stub promises a NestedDict.

Fix (type-only, apart from one bug)

• _compatible_imports.shaped_aval(atom): a typed accessor (a cast to ShapedArray). It is used at every flagged .aval.shape/.dtype/.ndim site.
• find_matched_vars / find_element_exist_in_the_set accept Var | Literal (they already filtered with isinstance).
• hid_param_op: trainable invars are Var | Literal. A literal never traces to a ParamState, and the loop narrows on that.
• Annotations or justified casts for:
  • the detach list;
  • the scan-descent body maps (hidden carry inits are already guarded to be Vars);
  • the module temporaries;
  • FlattedDict.split.
• new_var: # type: ignore[call-arg] on the pre-0.6.2 branch only.

Bug found on the way

IODimVjpAlgorithm.get_etrace_of checked etrace_x_key(relation.x_var) is not None. That value is id(x_var), 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. io_dim_vjp_test.py::TestStateLifecycle::test_get_etrace_of_elementwise_weight reproduces it: it fails with KeyError before the fix and passes after.

Verification

• python -m mypy braintrace: no issues.
• Full suite: 2987 passed.

Spec: docs/specs/2026-10-24-mypy-typecheck.md.

[codecov]

Codecov Report

❌ Patch coverage is 92.30769% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
braintrace/_compatible_imports.py	75.00%	1 Missing ⚠️
braintrace/_compiler/canonicalize.py	50.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!