Skip to content

CELL-20: event_backend="brainevent" uses brainevent.csrmv; "auto" resolves to scatter - #201

Merged
chaoming0625 merged 2 commits into
mainfrom
cell-20
Nov 9, 2026
Merged

chaoming0625 merged 2 commits into
mainfrom
cell-20

Conversation

@chaoming0625

@chaoming0625 chaoming0625 commented Nov 9, 2026 •

Copy link
Copy Markdown
Collaborator

CELL-20. event_backend="brainevent" had been unusable on every brainevent release since 0.1.0 because it called brainevent.coomv, which 0.1.0 removed. The tests that cover this path skipped when coomv was missing, so CI never ran it.

Decision

brainevent has a replacement with the same meaning: brainevent.csrmv(data, indices, indptr, v, *, shape, transpose, backend). Like coomv, it is a float sparse matrix-vector product, not the binary-gated variant. It has the same signature and the same jax_raw kernel (registered for cpu, gpu and tpu) in every release from 0.0.7 (the current floor) through 0.4.1. The brainevent path now uses it, and the brainevent>=0.0.7 floor stays as it is.

  • Each delay-grouped block is stored as a CSR matrix with one row per presynaptic member. Contacts are sorted stably by pre_index, and repeated destinations are kept. Weights are permuted when the op is called, so trainable weights still work. A weight with one element is passed as homogeneous.
  • The exact bilinear custom JVP built on scatter is unchanged.
  • "auto" now always resolves to "scatter". On brainevent >= 0.1 it already did. With the default brainevent_backend="jax_raw", the brainevent operator lowers to the same gather and scatter-add, so preferring it would change the default path for every user with no benefit.
  • If the installed brainevent has no csrmv, "brainevent" raises a RuntimeError during run/prepare_run setup. The message names event_backend='scatter'. make_delivery_op no longer falls back to scatter silently.

Changes

  • braincell/network/delivery.py: resolve_event_backend and make_delivery_op as described above.
  • delivery_test.py: the brainevent derivative test no longer skips. It spies on brainevent.csrmv to check that the path is really taken, and it uses unsorted pre_index, repeated destinations and a presynaptic member with no contacts. It compares primal, forward-mode, vmap-batched JVP, reverse-mode and bool-event results with scatter (rtol=1e-6). Backend-selection tests cover "auto"/"scatter"/"brainevent", and a brainevent without csrmv must raise an error that mentions scatter.
  • engine_test.py: the brainevent gradient test no longer skips. Two tests are new. The first runs scatter, brainevent and auto end to end; brainevent must match scatter with rtol=1e-6, and auto must equal scatter exactly. The second checks that prepare_run(event_backend="brainevent") is rejected with a "scatter" message when csrmv is missing.
  • CI.yml: the slow_tests deselect of bidirectional_test.py::...test_all_parameter_gradients_across_losses_delays_and_backends and its comment are removed.
  • changelog.md: a Bug Fixes entry.

Tests

  • Before the fix, the new and unskipped tests fail: 5 tests, 9 failures counting subtests. After the fix, pytest braincell/network/delivery_test.py braincell/network/engine_test.py gives 21 passed, 4 subtests passed. The same result holds with JAX_ENABLE_X64=1, and with brainevent 0.0.7, 0.1.0, 0.1.2, 0.2.1, 0.3.0, 0.4.0 and 0.4.1.
  • pytest "validation/optim/gradient_correctness/bidirectional_test.py::BidirectionalTest::test_all_parameter_gradients_across_losses_delays_and_backends" gives 1 passed, 4 subtests passed. pytest validation/optim -m slow gives 17 passed.
  • pytest braincell/ gives 2739 passed and 4 failed. The 4 failures are vis topology tests that need networkx, which is not installed in the local env. They are unrelated to this change.
  • pre-commit run --all-files passes.

Edge cases

@chaoming0625
chaoming0625 marked this pull request as ready for review November 9, 2026 00:44
@chaoming0625
chaoming0625 requested a review from TLDSZ as a code owner November 9, 2026 00:44
@chaoming0625
chaoming0625 merged commit 04ee37b into main Nov 9, 2026
16 checks passed
@chaoming0625
chaoming0625 deleted the cell-20 branch November 9, 2026 02:19
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant