Fix warning visibility and Map-based state initialization

[poilsosart]

Description

This PR updates warning handling and mapped-state initialization.

Changes

• Removed all executable warnings.catch_warnings and warnings.filterwarnings calls.
• Kept BrainTrace warnings visible to users.
• Replaced brainstate.transform.vmap_new_states with brainstate.nn.Map.
• Initialized mapped models with model.init_all_states().
• Fixed mapped parameter paths so model parameters, gradients, and optimizer keys match.
• Updated sequence drivers and examples for the new Map behavior.
• Replaced stale model.module.show_graph() calls with model.show_graph().

No algorithm logic or public function signatures were changed.

Testing

• Map regression tests: 3 passed.
• Focused compiler, sequence, Map, and convolution tests: 283 passed.
• Strict Sphinx documentation build passed.
• Full test suite: 2595 passed.

Remaining failures are caused by the existing saiunit.math.exprel(mV) compatibility issue and test-order contamination. These issues are outside the scope of this PR.

Summary by Sourcery

Adopt brainstate.nn.Map as the standard batching mechanism for braintrace, replacing vmap_new_states-based flows while keeping learner APIs intact, and remove warning suppression so compiler and algorithm diagnostics remain visible across code, tests, and documentation.

New Features:

• Switch Map-based batching to use brainstate.nn.Map for per-sample state ownership while keeping algorithm APIs unchanged.

Bug Fixes:

• Ensure mapped gradients and windowed sequence drivers preserve original model parameter paths and match sums over independent lanes.
• Preserve etp_conv primitives under mapped execution by registering a custom batching rule for brainstate.nn.Map.
• Deduplicate mapped hidden-state aliases so compiler paths remain stable and unique under Map.

Enhancements:

• Make braintrace.compile(..., vmap=True) return the concrete ETraceAlgorithm with a mapped model and expose reports and sequence drivers directly.
• Update sequence window-mode semantics so Map-based learners support chunked gradients and evolutions while legacy ETraceVmap wrappers remain guarded.
• Modernize examples, tests, and tutorials to use Map-based batching, direct learner.show_graph(), and unsuppressed diagnostics.
• Tighten warning handling by removing warning filters and relying on structured diagnostics instead of suppressing UserWarnings.
• Document the new warning-visibility and Map-initialization behavior in a dedicated spec.

Documentation:

• Rewrite batching, hidden-state, quickstart, and online-learning tutorials to describe Map-based batching and updated compile contracts.
• Refresh pp_prop and SNN tutorials to show Map-backed learners and unsuppressed compiler diagnostics.
• Update pp_prop example README and narrative text to reflect Map-based batching.

Tests:

• Extend sequence, convolution, compiler, and legacy API tests to cover Map-based gradients, window modes, mapped conv behavior, and deprecation warnings.
• Add regression tests ensuring Map initialization exposes clean state paths and that mapped compile results retain show_graph and reporting APIs.
• Adjust existing tests to stop suppressing warnings and to validate diagnostic behavior under control flow.

[sourcery-ai]

Sorry @poilsosart, you have reached your weekly rate limit of 500000 diff characters.

Please try again later or upgrade to continue using Sourcery

[sourcery-ai]

Reviewer's Guide

Replaces vmap-based batching with Map-based state initialization across compile paths, tests, examples, and docs while removing warning suppression so BrainTrace diagnostics surface to users, and ensuring mapped models preserve parameter paths and ETP primitive behavior.

Sequence diagram for braintrace.compile(..., vmap=True) with Map-based initialization

sequenceDiagram
    actor User
    participant braintrace_compile as braintrace.compile
    participant Map as brainstate_nn_Map
    participant Algo as ETraceAlgorithm

    User->>braintrace_compile: compile(model, algo_cls, example_inputs, batch_size=B, vmap=True)
    braintrace_compile->>Map: Map(model, init_map_size=B)
    braintrace_compile->>Map: init_all_states()
    braintrace_compile->>Algo: Algo(Map, **options)
    braintrace_compile->>Algo: compile_graph(*example_inputs)
    braintrace_compile-->>User: Algo

Loading

File-Level Changes

Change	Details	Files
Switch compile(vmap=True) from vmap_new_states + Vmap wrapper to brainstate.nn.Map with init_all_states and direct algorithm return.	Updated braintrace.compile to wrap models in brainstate.nn.Map(model, init_map_size=batch_size) when vmap=True and call model.init_all_states() before compiling with batched example inputs. Removed the ETraceVmap wrapper from the compile path so compile always returns an ETraceAlgorithm, while retaining ETraceVmap only as a legacy helper for explicitly constructed brainstate.nn.Vmap models. Adjusted sequence drivers (SequenceDriverMixin) and window-mode checks to treat only legacy ETraceVmap wrappers as vmapped, allowing windowed etrace_grad/etrace_evolve under Map-backed learners. Extended compile tests to assert Map ownership of graph_executor.model, preserved map size, and direct access to report/show_graph on the learner.	braintrace/_compile.py braintrace/_algorithm/sequence.py braintrace/_algorithm/base.py braintrace/_algorithm/tests/compile_test.py examples/tests/test_compile_modes.py
Ensure Map-backed learners preserve parameter/state paths and ETP primitive behavior, including conv batching under mapped execution.	Added a custom batching rule for etp_conv that preserves the ETP primitive when only the convolution input is mapped, falling back to default batching otherwise. Recorded ETP primitive batching rules in a shared _ETP_BATCHING_RULES registry for reuse by custom batchers. Adjusted module-info extraction to deduplicate Map state aliases, preferring module-facing paths over internal dict_vmap_states entries, and to special-case brainstate.nn.Map so compiler state discovery uses the wrapped module. Introduced tests verifying Map state-path deduplication, conv Map gradients equal sums over single-sample grads, and stable behavior for mixed conv/dense and LayerNorm pipelines under Map. Updated conv correctness tests and docstrings from Vmap/vmap_new_states terminology to Map and mapped execution semantics.	braintrace/_op/conv.py braintrace/_op/_primitive.py braintrace/_compiler/module_info.py braintrace/_compiler/module_info_test.py braintrace/_algorithm/tests/conv_vmap_correctness_test.py
Replace vmap-based sequence tests with Map-backed tests and strengthen gradient, window, and parameter-path invariants.	Renamed vmap-specific helpers and tests to Map variants (e.g., _vmap_learner to _map_learner, TestVmap to TestMap) and updated expectations to assert learners are ETraceAlgorithm instances with brainstate.nn.Map models rather than brainstate.nn.Vmap. Added tests to ensure gradients under Map keep original model ParamState paths, optimizer keys align with learner.param_states, and lane-wise gradients sum to mapped gradients for both plain and windowed modes. Extended sequence tests to cover window-mode gradients and etrace_evolve outputs under Map, including comparisons between mapped window gradients and sums over independently driven lane windows, and to verify chunk_size=1 behaves as the plain path. Kept a legacy vmap compatibility test ensuring ETraceVmap still exposes sequence drivers while clarifying its status as a compatibility type only for manually constructed Vmap learners.	braintrace/_algorithm/sequence_test.py
Remove warning suppression so BrainTrace diagnostics are visible in tests, examples, and docs, updating tests to assert on warnings via pytest when needed.	Removed all uses of warnings.catch_warnings and warnings.filterwarnings from compiler, oracle, hidden-group, scenario, canonicalization, while-support, diagnostic-exploration, and concept tests, letting diagnostics surface while still asserting on structured DiagnosticKind records. Updated deprecation tests for legacy ops/params to use pytest.warns and recwarn fixtures instead of manual warnings.catch_warnings, explicitly asserting presence or absence of DeprecationWarning as appropriate. Adjusted control-flow and scan-descent tests to assert absence of UserWarnings via recwarn rather than suppressing them, and ensured downgrade of some diagnostics to INFO is reflected only in structured records. Stopped filtering expected ETP primitive decomposition warnings in conv-related tests, relying instead on the new Map-aware conv batching rule to avoid noisy decomposition paths in vmapped scenarios.	braintrace/_compiler/scenario_catalog_test.py braintrace/_compiler/tests/compiler_property_test.py braintrace/_compiler/tests/compiler_oracle_test.py braintrace/_compiler/hidden_group_test.py braintrace/_compiler/canonicalize_test.py braintrace/_compiler/tests/cell_relation_guardrail_test.py braintrace/_algorithm/tests/diagnostic_exploration_test.py braintrace/_algorithm/tests/while_support_test.py braintrace/_legacy/_ops_test.py braintrace/_legacy/_params_test.py
Migrate examples, tutorials, and docs from vmap_new_states/Vmap to Map-based batching and remove warning filters from public docs.	Rewrote documentation notebooks (batching, hidden_states, RNN/SNN online learning, DRTRL, pp_prop, neural_network_layers, quickstart) to describe Map-based batching, including updated explanations of compile(..., vmap=True) semantics, Map-based state initialization, and direct learner usage. Updated code cells to construct brainstate.nn.Map(model, init_map_size=B), call mapped_model.init_all_states(), and compile from batched inputs; removed use of brainstate.transform.vmap_new_states and brainstate.nn.Vmap wrappers in examples and tutorials. Simplified doc examples by removing warning filters around compiler calls, instead noting in comments when non-temporal readouts emit diagnostics, and letting Sphinx builds fail on unexpected warnings. Aligned pp_prop README and narrative examples (including 05-batching-vmap.py, 02-batching-vmap.py, and conv/feedforward SNN examples) with the Map workflow and updated show_graph usage from model.module.show_graph() to learner.show_graph().	docs/advanced/batching.ipynb docs/tutorials/hidden_states.ipynb docs/tutorials/rnn_online_learning.ipynb docs/tutorials/snn_online_learning.ipynb docs/tutorials/drtrl.ipynb docs/tutorials/pp_prop.ipynb docs/tutorials/neural_network_layers.ipynb docs/quickstart/quickstart.ipynb docs/quickstart/concepts.ipynb examples/004-feedforward-conv-snn.py examples/drtrl/02-batching-vmap.py examples/pp_prop/05-batching-vmap.py examples/pp_prop/_shared.py examples/snn_models.py examples/002-coba-ei-rsnn.py examples/100-gru-on-copying-task.py examples/pp_prop/12-classification-neuromorphic.py examples/pp_prop/14-knob-vjp-method-contrast.py examples/tests/test_compile_modes.py examples/003-snn-memory-and-speed-evaluation-vmap.py examples/003-snn-memory-and-speed-evaluation-batched.py examples/003-snn-memory-and-speed-evaluation-all.py docs/specs/2026-07-28-warnings-and-map-initialization.md

---

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!