fix(brainpy): repair 18 correctness bugs found in a full-package audit#868
Merged
Conversation
A systematic audit of the brainpy/ package surfaced 18 confirmed defects across dnn, integrators, encoding, initialize, connect, optim, dynold, algorithms, math, runners and train. Each fix ships with a co-located regression test (sibling *_test.py). (The adaptive-pooling upsampling crash surfaced by the same audit was fixed independently on master in #867, so it is not duplicated here.) Critical - dnn/normalization: LayerNorm reduced over the leading batch/time axes instead of the trailing normalized_shape dimensions. - integrators/sde/normal: Stratonovich Euler-Heun predictor used the raw standard-normal noise instead of the sqrt(dt)-scaled Wiener increment, making it inconsistent with the corrector. - encoding/stateful_encoding: WeightedPhaseEncoder used an integer power base 2 ** (-k), which truncates to 0 under int tracing and collapsed every phase weight to zero (the encoder never spiked). High - initialize/random_inits: TruncatedNormal(lower=None, upper=None) raised TypeError; default to 2-sigma truncation, treat None as +/- inf. - initialize/decay_inits: DOGDecay produced an asymmetric connectivity matrix on non-square grids (bad meshgrid moveaxis) and never applied its normalize flag. - connect/random_conn: FixedProb.build_csr built a structurally-invalid CSR indptr when pre_ratio < 1; FixedTotalNum raised TypeError on a fractional num. - optim/optimizer: MomentumNesterov was byte-identical to plain Momentum (missing the Nesterov look-ahead step). - dynold (synplast + experimental) STP: the continuous facilitation ODE carried a spurious +U source term, so u drifted toward U*tau_f even with no presynaptic spikes. - algorithms/offline: the Ridge/Regression IRLS while_loop seeded par_old within allclose tolerance of par_new, so it exited immediately and returned the untrained initial weights. Medium - connect/regular_conn: GridConn periodic boundary produced duplicate (pre, post) COO edges on small grids (e.g. 2x2 GridFour double-counted). - math/object_transform/autograd: vector_grad used as a parameterized decorator (func=None) crashed instead of returning a decorator. - optim/optimizer: Adadelta silently ignored its learning rate. - runners: DSRunner.predict dropped shared_args in the memory-efficient path. - algorithms/offline: LogisticRegression Newton path built its diagonal via in-place bm.Array assignment, tripping the JAX __jax_array__-during- abstractification error under while_loop. - train/online: OnlineTrainer._fit transposed histories unconditionally, crashing on scalar monitor histories and returning batch-major output in data_first_axis='T' mode. Low - initialize/random_inits: _compute_fans raised IndexError on 0-D/1-D shapes, breaking every VarianceScaling init on bias vectors. - optim/optimizer: LARS zero-norm guard used maximum (could exceed the trust ratio); switched to where to clamp to exactly 1.0. - train/back_propagation: fit rejected a valid iterable dataset holding exactly two (x, y) batches as if it were the removed bare (X, Y) API.
…atio<1 Two coverage tests pinned the pre-fix malformed CSR (a truncated indptr of length int(pre_num*pre_ratio)+1). Update them to assert the correct CSR shape (pre_num+1 entries, non-selected pre rows have zero out-degree) that the H4 fix now produces, plus internal indptr/indices consistency.
Reviewer's GuideThis PR fixes 18 audited correctness bugs across optimizers, initialization, connectivity, normalization, runners, training, encoders, SDE integrators, STP dynamics, autograd, and offline algorithms, with co-located regression tests and some legacy test cleanup. Sequence diagram for DSRunner.predict shared_args propagationsequenceDiagram
actor User
participant DSRunner
participant _fun_predict
participant run_fun_step
User->>DSRunner: predict(inputs, shared_args)
DSRunner->>_fun_predict: _fun_predict(indices, *inputs, shared_args)
_fun_predict->>_fun_predict: choose _jit_step_func_predict or _step_func_predict
_fun_predict->>run_fun_step: functools.partial(run_fun_step, shared_args=shared_args)
loop over time indices
_fun_predict->>run_fun_step: run_fun_step(i, *inputs, shared_args)
run_fun_step-->>_fun_predict: step_outputs
end
_fun_predict-->>DSRunner: stacked step_outputs
DSRunner-->>User: predictions
Sequence diagram for Stratonovich Euler-Heun SDE step noise scalingsequenceDiagram
participant SDEIntegrator as NormalSDEIntegrator
participant Random as bm_random
participant Diffusion as g
SDEIntegrator->>SDEIntegrator: step(*args, **kwargs)
SDEIntegrator->>bm_random: randn(*shape)
bm_random-->>SDEIntegrator: noise
SDEIntegrator->>SDEIntegrator: all_noises[key] = noise * sqrt(dt)
alt VECTOR_WIENER
SDEIntegrator->>Diffusion: diffusions[key]
Diffusion-->>SDEIntegrator: diffusions[key]
SDEIntegrator->>SDEIntegrator: y_bar = all_args[key] + sum(diffusions[key] * all_noises[key])
else
SDEIntegrator->>Diffusion: diffusions[key]
Diffusion-->>SDEIntegrator: diffusions[key]
SDEIntegrator->>SDEIntegrator: y_bar = all_args[key] + diffusions[key] * all_noises[key]
end
SDEIntegrator->>Diffusion: diffusion_bars = g(**all_args_bar)
Diffusion-->>SDEIntegrator: diffusion_bars
SDEIntegrator-->>SDEIntegrator: corrector step using diffusion_bars
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - I've left some high level feedback:
- The periodic-boundary de-duplication in
GridConn.build_coo()usesnp.uniqueon the full(pre, post)pair array, which could be quite expensive on large grids; consider a more incremental or JAX-native de-duplication strategy to avoid the extra host round-trip and quadratic-ish behavior on big connectivity sets. - The new
FixedTotalNumhandling for floatnumroundsnum * mat_element_numto an integer; if callers expect a strict fraction or a maximum count, you may want to clamp (e.g., floor and enforce a minimum of 1) or explicitly document the rounding semantics to avoid surprising over- or under-selection on edge cases.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The periodic-boundary de-duplication in `GridConn.build_coo()` uses `np.unique` on the full `(pre, post)` pair array, which could be quite expensive on large grids; consider a more incremental or JAX-native de-duplication strategy to avoid the extra host round-trip and quadratic-ish behavior on big connectivity sets.
- The new `FixedTotalNum` handling for float `num` rounds `num * mat_element_num` to an integer; if callers expect a strict fraction or a maximum count, you may want to clamp (e.g., floor and enforce a minimum of 1) or explicitly document the rounding semantics to avoid surprising over- or under-selection on edge cases.Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Merged
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
1 participant
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.
Summary
A systematic audit of the whole
brainpy/package surfaced 18 confirmedcorrectness defects, spanning 14 source modules. This PR fixes all of them,
each with a co-located regression test (sibling
*_test.py, TDD-style).Severity breakdown: 3 Critical, 6 High, 6 Medium, 3 Low.
Highlights
trailing
normalized_shape).sqrt(dt).ignored its learning rate.
crashed under current JAX.
Full list
Critical
dnn/normalization: LayerNorm reduced over leading batch/time axes instead of the trailingnormalized_shape.integrators/sde/normal: Stratonovich Euler-Heun predictor used raw noise instead of thesqrt(dt)-scaled Wiener increment.encoding/stateful_encoding: WeightedPhaseEncoder integer power base collapsed every phase weight to 0 (never spiked).High
initialize/random_inits:TruncatedNormal(lower=None, upper=None)raisedTypeError; default to 2-sigma, treatNoneas ±inf.initialize/decay_inits:DOGDecayproduced an asymmetric matrix on non-square grids and never applied itsnormalizeflag.connect/random_conn:FixedProb.build_csrbuilt a malformed CSR indptr forpre_ratio<1;FixedTotalNumraisedTypeErroron a fractionalnum.optim/optimizer:MomentumNesterovwas byte-identical to plain Momentum (no look-ahead).dynoldSTP (synplast + experimental): the facilitation ODE carried a spurious+Usource term, soudrifted with no spikes.algorithms/offline: the Ridge/Regression IRLSwhile_loopseeded within tolerance and returned untrained weights.Medium
connect/regular_conn:GridConnperiodic boundary produced duplicate COO edges on small grids.math/object_transform/autograd:vector_gradparameterized-decorator form (func=None) crashed.optim/optimizer:Adadeltasilently ignored its learning rate.runners:DSRunner.predictdroppedshared_argsin the memory-efficient path.algorithms/offline:LogisticRegressionNewton path tripped the JAX__jax_array__-during-abstractification error underwhile_loop.train/online:OnlineTrainer._fittransposed histories unconditionally — crashed on scalar monitors, wrong layout indata_first_axis='T'.Low
initialize/random_inits:_compute_fansraisedIndexErroron 0-D/1-D shapes (broke VarianceScaling on bias vectors).optim/optimizer:LARSzero-norm guard usedmaximum(could exceed trust ratio); switched towhere.train/back_propagation:fitrejected a valid iterable of two(x, y)batches as the removed bare(X, Y)API.Also included
This branch was reconciled with local
master, which additionally removes threelegacy, superseded test files (
brainpy/math/numpy_{einsum,indexing,ops}_test.py).They appear in this diff because they were not yet on the remote
master.Testing
brainpy/suite green (4248 passed, 0 failed, 12 skipped).Summary by Sourcery
Fix multiple correctness issues across optimization, initialization, normalization, connectivity, encoding, integrators, training, runners, autograd, and offline learning components in the brainpy package, adding regression tests for each bug and cleaning up legacy numpy-based tests.
Bug Fixes:
Tests: