Migrate sparse_matmul off brainevent's deprecated yw_to_w protocol

[chaoming0625]

Summary

• brainevent 0.1.2 renames its DataRepresentation trace protocol from
  yw_to_w/yw_to_w_transposed to dt2t/dt2t_transposed (mirroring
  braintrace's own DT_TO_T rename in Rename YW_TO_W ETP rule to DT_TO_T #130), keeping the old names as
  deprecated aliases slated for removal. braintrace._op.sparse now calls
  .dt2t/.dt2t_transposed directly instead of the deprecated aliases, and
  the test stubs (_StubSparseMat, _HashableStub, _DuckMat) are updated to
  match. Minimum brainevent is raised to 0.1.2 in pyproject.toml /
  requirements.txt.
• Fixes a pre-existing JAX-internal linear_util.wrap_init
DeprecationWarning in the single-step VJP residual construction
  (vjp_graph_executor.py), by threading a DebugInfo object through the
  ecosystem-standard brainstate._compatible_import.wrap_init shim
  (re-exported as braintrace._compatible_imports.wrap_init), matching the
  pattern already used elsewhere in the brainstate/saiunit stack.
• One doc-comment fix in docs/tutorials/etp_primitives.ipynb (brainevent
  protocol reference only — the notebook's stale ETP_RULES_YW_TO_W references
  to braintrace's own pre-Rename YW_TO_W ETP rule to DT_TO_T #130 rule name are a separate, pre-existing issue
  and out of scope here).

No behavior change; both migrations are mechanical/annotation-only.

Test plan

• pytest braintrace/ — 2123 passed, 0 failed, 3 skipped (matches baseline)
• pytest braintrace/_op/sparse_test.py braintrace/nn/_linear_test.py -W error::DeprecationWarning — 107 passed, zero brainevent/JAX deprecation warnings surfaced
• mypy braintrace — no issues found

Summary by Sourcery

Migrate sparse matmul and related tests to brainevent’s updated dt2t/dt2t_transposed trace protocol and address a JAX wrap_init deprecation in the VJP graph executor, with corresponding documentation and dependency updates.

Enhancements:

• Update sparse_matmul implementation and eligibility-trace rules to use brainevent.DataRepresentation.dt2t and dt2t_transposed instead of the deprecated yw_to_w variants.
• Adjust sparse-matrix stubs and protocol checks in sparse tests and linear layer tests to align with the new brainevent trace-propagation API.
• Thread DebugInfo via the ecosystem-standard wrap_init shim when constructing single-step VJP residuals to silence JAX linear_util.wrap_init deprecation warnings.
• Re-export wrap_init from brainstate._compatible_import through braintrace._compatible_imports to provide a unified compatibility layer.

Build:

• Raise the minimum brainevent dependency to version 0.1.2 in pyproject.toml and requirements.txt.

Documentation:

• Refresh ETP primitives tutorial text to describe the brainevent DataRepresentation with_data / dt2t_transposed / dt2t protocol.
• Add unreleased changelog entry describing the sparse_matmul brainevent protocol migration, the JAX deprecation fix, and the raised brainevent minimum version.

Tests:

• Update sparse and SparseLinear test suites, including stubs and error messages, to target the dt2t/dt2t_transposed protocol and ensure compatibility with brainevent>=0.1.2.

[sourcery-ai]

Reviewer's Guide

Migrates braintrace’s sparse matmul integration from brainevent’s deprecated yw_to_w* protocol to the new dt2t* protocol, updates stubs/tests/docs accordingly, bumps the brainevent minimum version, and fixes a JAX linear_util.wrap_init deprecation warning by routing VJP residual construction through the ecosystem wrap_init shim.

Sequence diagram for sparse_matmul using brainevent DataRepresentation.dt2t_transposed

sequenceDiagram
    actor User
    participant sparse_matmul
    participant _sp_mm_dt_to_t
    participant jax_vmap as jax.vmap
    participant DataRepresentation

    User->>sparse_matmul: sparse_matmul(x, weight_data, sparse_mat)
    sparse_matmul->>_sp_mm_dt_to_t: _sp_mm_dt_to_t(hidden_dim, trace, sparse_mat)
    _sp_mm_dt_to_t->>DataRepresentation: _unwrap_sparse_mat(sparse_mat)
    alt hidden_dim.ndim == 2
        _sp_mm_dt_to_t->>jax_vmap: jax.vmap(DataRepresentation.dt2t_transposed)
        jax_vmap->>DataRepresentation: dt2t_transposed(hidden_dim_batch, weight_trace_batch)
    else
        _sp_mm_dt_to_t->>DataRepresentation: dt2t_transposed(hidden_dim, weight_trace)
    end
    _sp_mm_dt_to_t-->>sparse_matmul: out['weight']
    sparse_matmul-->>User: y_sp

Loading

Sequence diagram for VJP residual construction using wrap_init shim

sequenceDiagram
    participant call_over_single_step
    participant wrap_init
    participant flatten_fun as jax.api_util.flatten_fun_nokwargs
    participant trace_jaxpr as pe.trace_to_jaxpr_dynamic

    call_over_single_step->>call_over_single_step: vjp_cotangent_args = ((out_single_or_multi_steps, etrace_state_vals, other_state_vals),)
    call_over_single_step->>call_over_single_step: out_flat, out_tree = jax.tree.flatten(vjp_cotangent_args)
    call_over_single_step->>wrap_init: wrap_init(f_vjp, vjp_cotangent_args, {}, braintrace_vjp_residual)
    wrap_init-->>call_over_single_step: wrapped_f_vjp
    call_over_single_step->>flatten_fun: flatten_fun_nokwargs(wrapped_f_vjp, out_tree)
    flatten_fun-->>call_over_single_step: rule, in_tree
    call_over_single_step->>trace_jaxpr: trace_to_jaxpr_dynamic(rule, out_avals)
    trace_jaxpr-->>call_over_single_step: jaxpr, consts
    call_over_single_step->>call_over_single_step: residual = VjpResiduals(jaxpr, in_tree(), out_tree, consts)

Loading

File-Level Changes

Change	Details	Files
Update sparse matmul implementation and tests to use brainevent.DataRepresentation.dt2t / dt2t_transposed instead of deprecated yw_to_w / yw_to_w_transposed, and tighten brainevent version requirements.	Change sparse_matmul and internal ETP rules to call mat.dt2t_transposed (and mat.dt2t conceptually) instead of mat.yw_to_w_transposed / mat.yw_to_w. Adjust sparse DataRepresentation stubs in tests (_StubSparseMat, _HashableStub, _DuckMat) to implement dt2t_transposed / dt2t instead of yw_to_w_transposed / yw_to_w. Update error messages and test expectations to reference dt2t_transposed / dt2t in type checks and protocol documentation. Refresh tutorial notebook text to describe the with_data / dt2t_transposed / dt2t protocol for brainevent.DataRepresentation. Raise brainevent dependency minimum to 0.1.2 in pyproject.toml and requirements.txt.	braintrace/_op/sparse.py braintrace/_op/sparse_test.py braintrace/nn/_linear_test.py docs/tutorials/etp_primitives.ipynb pyproject.toml requirements.txt
Fix JAX DeprecationWarning around linear_util.wrap_init in VJP residual construction by using the shared brainstate/braintrace compatible wrap_init shim.	Import wrap_init from braintrace._compatible_imports and re-export it from brainstate._compatible_import._compatible_imports for downstream use. Replace direct lu.wrap_init(f_vjp) usage with wrap_init(f_vjp, vjp_cotangent_args, {}, 'braintrace_vjp_residual') to supply DebugInfo and avoid deprecated API paths. Thread the cotangent argument tuple (out_single_or_multi_steps, etrace_state_vals, other_state_vals) consistently into flatten_fun_nokwargs via the new wrap_init call.	braintrace/_algorithm/vjp_graph_executor.py braintrace/_compatible_imports.py changelog.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.

[sourcery-ai]

Hey - I've reviewed your changes and they look great!

---

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.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!