Repository navigation
CELL-20: event_backend="brainevent" uses brainevent.csrmv; "auto" resolves to scatter - #201
Merged
Merged
Conversation
chaoming0625
marked this pull request as ready for review
November 9, 2026 00:44
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
CELL-20.
event_backend="brainevent"had been unusable on every brainevent release since 0.1.0 because it calledbrainevent.coomv, which 0.1.0 removed. The tests that cover this path skipped whencoomvwas missing, so CI never ran it.Decision
brainevent has a replacement with the same meaning:
brainevent.csrmv(data, indices, indptr, v, *, shape, transpose, backend). Likecoomv, it is a float sparse matrix-vector product, not the binary-gated variant. It has the same signature and the samejax_rawkernel (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 thebrainevent>=0.0.7floor stays as it is.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."auto"now always resolves to"scatter". On brainevent >= 0.1 it already did. With the defaultbrainevent_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.csrmv,"brainevent"raises aRuntimeErrorduringrun/prepare_runsetup. The message namesevent_backend='scatter'.make_delivery_opno longer falls back to scatter silently.Changes
braincell/network/delivery.py:resolve_event_backendandmake_delivery_opas described above.delivery_test.py: the brainevent derivative test no longer skips. It spies onbrainevent.csrmvto check that the path is really taken, and it uses unsortedpre_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 withoutcsrmvmust 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 withrtol=1e-6, and auto must equal scatter exactly. The second checks thatprepare_run(event_backend="brainevent")is rejected with a "scatter" message whencsrmvis missing.CI.yml: theslow_testsdeselect ofbidirectional_test.py::...test_all_parameter_gradients_across_losses_delays_and_backendsand its comment are removed.changelog.md: a Bug Fixes entry.Tests
pytest braincell/network/delivery_test.py braincell/network/engine_test.pygives 21 passed, 4 subtests passed. The same result holds withJAX_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 slowgives 17 passed.pytest braincell/gives 2739 passed and 4 failed. The 4 failures are vis topology tests that neednetworkx, which is not installed in the local env. They are unrelated to this change.pre-commit run --all-filespasses.Edge cases
JAX_ENABLE_X64=1.csrmvis covered by the delivery and engine tests.ImportErrorbranch, which the changed-lines coverage gate had flagged (test_event_backend_selection_without_brainevent).slow_testsincluded (run 42541898789), so the bidirectional test now passes in CI without the deselect.ci-okpasses on 187f2f4.