fix: restore daily CI under pytest 9.1 + brainstate 0.5 typed API

[chaoming0625]

Summary

The scheduled daily CI began failing on the frozen commit d0627d4 purely from upstream dependency drift (CI was green 2026-06-01, red from 2026-06-02 on the same commit). Diagnosis found four independent root causes; this PR fixes the three braintrace-side ones. The fourth was an upstream saiunit bug, now fixed in saiunit/brainunit 0.5.1 — no braintrace change required.

#	Failure	Root cause	Fix
1	Collection error (GraphNodeMeta has no len())	pytest 9.1.0 parses 'cls,' trailing comma as 2 values	drop trailing commas in 4 compiler test files
2	154 mypy errors	brainstate 0.4+ ships py.typed	adopt the typed API in source
3	5 conv tests	brainstate 0.5.0 hardened validation (assert→ValueError) + padding-tuple semantics	update _conv_test expectations
4	40 SNN VJP tests	saiunit 0.4.0 str-mantissa guard crashed JAX 0.10.1's custom_vjp error-rendering for any unit-carrying output	fixed upstream in saiunit 0.5.1

Details

• New version #1 — pytest 9.1.0's ParameterSet._for_parametrize length check treats a single-arg id with a trailing comma ('cls,') as two parameters. Removed the trailing commas (module_info_test, hidden_group_test, hid_param_op_test, hidden_pertubation_test).
• Add etrace learning algorithm with O(n^2) complexity, and hybrid O(n) & O(n^2) algorithms #2 — Routed PyTree through braintrace's existing alias, centralized an as_size_tuple() helper in _typing, dropped FlattedDict subscripts, and added boundary asserts/casts. # type: ignore only where brainstate's typing makes it unavoidable (hidden Module.__init__, saiunit Unit|Quantity, treedef-as-PyTree).
• for pre- and post-eligibility trace, add independent decay factors and time constants #3 — brainstate 0.5.0 raises ValueError (not bare assert) for invalid groups/padding-string/in_size, and interprets an int padding tuple as one value per spatial dim. Updated the affected expectations only.
• algorithm updates #4 — No braintrace bug: JAX 0.10.1 eagerly builds a display tree via tree_unflatten(out_tree, [aval.str_short() …]) (string leaves) before the structure-match check; saiunit 0.4.0's new str-mantissa guard crashed on those strings for every custom_vjp carrying a Quantity. Fixed in saiunit by letting tree_unflatten round-trip placeholder leaves.

Verification (local, mirrors CI)

• pytest braintrace/ → 1367 passed, 3 skipped, 2 xfailed, 0 failed
• mypy braintrace → Success: no issues found in 51 source files
• python -m build → wheel + sdist build; py.typed shipped in both (PEP 561)

🤖 Generated with Claude Code

Summary by Sourcery

Restore compatibility with newer pytest and the typed brainstate/brainunit stack by tightening typing, normalizing size handling, and aligning tests with updated runtime validation and PyTree semantics.

Bug Fixes:

• Fix RNN, readout, and tracing components to use normalized size tuples and robust PyTree/unit handling so they work with brainstate 0.5's typed API.
• Adjust convolution tests and padding/group validation expectations to match brainstate's updated error types and padding semantics.
• Update parametrized compiler tests to avoid pytest 9.1 collection errors caused by single-argument parameter IDs with trailing commas.

Enhancements:

• Introduce a shared as_size_tuple helper and centralize PyTree aliases to satisfy static type checking with brainstate's exported typing, including safer use of FlattedDict and cached PyTreeDefs.
• Harden algorithm and executor internals with explicit type annotations, runtime asserts, and nullability checks for state splits and filtered state views.
• Guard docstring propagation and registry definitions to be type-checker friendly under the new dependency versions.

[sourcery-ai]

Reviewer's Guide

Restores daily CI under pytest 9.1 and brainstate 0.5 by tightening braintrace’s typing around brainstate’s new typed API, normalizing size handling in RNN/readout layers, relaxing FlattedDict/PyTree generics in ETrace internals, updating conv tests for changed validation semantics, and fixing pytest parametrize collections; includes small robustness tweaks (docstring guards, explicit casts/asserts) and a type-annotated algorithm registry.

File-Level Changes

Change	Details	Files
Normalize in_size/out_size handling via a shared helper and make RNN/readout layers type-safe under brainstate’s typed Size API.	Introduce as_size_tuple(Size) in _typing to normalize scalar/sequence sizes to tuple[int, ...]. Update RNN modules and LeakyRateReadout to store in_size/out_size via as_size_tuple and use normalized sizes when constructing Linear layers and random initializers. Add Any/dict[str, Any] annotations where needed and add targeted type: ignore comments for brainstate.Module.init and saiunit tau typing.	braintrace/_typing.py braintrace/nn/_rnn.py braintrace/nn/_readout.py
Align internal ETrace/grad logic with brainstate’s PyTree typing and relax FlattedDict generics to satisfy mypy.	Alias and import PyTree from braintrace._typing and use it instead of brainstate.typing.PyTree across state management, gradient accumulation, and vjp helpers. Change state mapping parameters from Dict to Mapping where appropriate, and adjust function signatures/return types accordingly. Drop FlattedDict[...] generics on properties/attributes, add Optional typing and asserts after lazy initialization, and assert isinstance(...) where brainstate returns a union of mapping types.	braintrace/_typing.py braintrace/_state_managment.py braintrace/_etrace_algorithms/base.py braintrace/_etrace_algorithms/_common.py braintrace/_grad_exponential.py braintrace/_etrace_algorithms/param_dim_vjp.py braintrace/_etrace_algorithms/graph_executor.py braintrace/_etrace_algorithms/oracle.py braintrace/_etrace_compiler/module_info.py braintrace/_etrace_compiler/hidden_group.py braintrace/_etrace_compiler/hid_param_op.py
Update conv tests and pytest parametrization to match upstream behavioral changes in brainstate and pytest 9.1.	Adjust Conv tests to expect ValueError instead of AssertionError for invalid groups/padding/in_size and to use explicit 3D padding tuples. Remove trailing commas from single-parameter pytest.mark.parametrize argument lists to avoid pytest 9.1 treating them as two parameters. Clarify conv padding docstrings to document per-spatial-dimension tuple semantics.	braintrace/nn/_conv_test.py braintrace/_etrace_compiler/hidden_group_test.py braintrace/_etrace_compiler/hid_param_op_test.py braintrace/_etrace_compiler/hidden_pertubation_test.py braintrace/_etrace_compiler/module_info_test.py
Harden small API surfaces and docs around Linear/Conv wrappers and algorithm registry for better typing and robustness.	Guard access to upstream brainstate docstrings by falling back to '' when doc is None in Linear/Conv wrapper classes before replacing the module name. Annotate the algorithm registry mapping with concrete key/value types. Add a type: ignore on ModuleInfo’s treedef.unflatten call to document the runtime PyTreeDef type mismatch with brainstate’s broad PyTree typing.	braintrace/nn/_linear.py braintrace/nn/_conv.py braintrace/_compile.py braintrace/_etrace_compiler/module_info.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've left some high level feedback:

• In the RNN modules, self.in_size/self.out_size are already normalized via _as_size_tuple in __init__, but many later calls (e.g. to Linear and _forget_bias) re-wrap these attributes with _as_size_tuple on every use; consider normalizing once and either using the attributes directly or caching their last dimension in a local variable to reduce repetition and improve readability.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In the RNN modules, `self.in_size`/`self.out_size` are already normalized via `_as_size_tuple` in `__init__`, but many later calls (e.g. to `Linear` and `_forget_bias`) re-wrap these attributes with `_as_size_tuple` on every use; consider normalizing once and either using the attributes directly or caching their last dimension in a local variable to reduce repetition and improve readability.

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}