fix(brainpy): repair 18 correctness bugs found in a full-package audit (#868)

A second audited, library-wide bug-fix sweep: 18 confirmed defects across dnn, integrators, encoding, initializers, connectivity, optimizers, dynold STP, offline/online training, object transforms, and runners — each with a co-located regression test. Also removes three legacy numpy_*_test.py files carried over from local master.

Author: chaoming0625

brainpy/algorithms/offline.py

@@ -173,8 +173,12 @@ def body_fun(a):
             par_new2 = par_new - self.learning_rate * grad_w
             return i + 1, par_new, par_new2
 
-        # Tune parameters for n iterations
-        r = while_loop(cond_fun, body_fun, (0, w - 1e-8, w))
+        # Tune parameters for n iterations. Seed ``par_old`` FAR from ``par_new`` so
+        # the very first convergence check (``not allclose(par_old, par_new)``) passes
+        # and the body runs at least once. ``w - 1e-8`` is within ``allclose``'s default
+        # tolerance of ``w``, so the loop exited immediately and returned the untrained
+        # initial weights (H7, audit 2026-07-08). Matches the ridge path's ``w + 1.``.
+        r = while_loop(cond_fun, body_fun, (0, w + 1., w))
         return r[-1]
 
     def predict(self, W, X):
@@ -415,10 +419,14 @@ def body_fun(a):
                 # respect to the parameters to minimize the loss
                 par_new2 = par_new - self.learning_rate * (y_pred - targets).dot(inputs)
             else:
-                gradient = self.sigmoid.grad(inputs.dot(par_new))
-                diag_grad = bm.zeros((gradient.size, gradient.size))
-                diag = bm.arange(gradient.size)
-                diag_grad[diag, diag] = gradient
+                # Build the diagonal weight matrix with pure JAX ops. The previous
+                # ``bm.zeros(...)`` + in-place ``diag_grad[diag, diag] = ...`` created a
+                # ``brainpy.Array`` whose ``__jax_array__`` protocol is invoked during
+                # ``while_loop`` abstractification, which current JAX rejects
+                # ("Triggering __jax_array__() during abstractification is no longer
+                # supported") (M5, audit 2026-07-08).
+                gradient = bm.as_jax(self.sigmoid.grad(inputs.dot(par_new)))
+                diag_grad = jnp.diag(gradient)
                 par_new2 = jnp.linalg.pinv(inputs.T.dot(diag_grad).dot(inputs)).dot(inputs.T).dot(
                     diag_grad.dot(inputs).dot(par_new) + targets - y_pred)
             return i + 1, par_new, par_new2

brainpy/algorithms/offline_test.py

@@ -91,6 +91,16 @@ def test_gradient_descent_path(self):
         w = np.asarray(bm.as_jax(algo(y, x)))
         assert np.all(np.isfinite(w))
 
+    def test_gradient_descent_actually_fits(self):
+        # H7 regression: the while_loop was seeded with ``(0, w - 1e-8, w)`` and
+        # ``allclose(w - 1e-8, w)`` is True, so the convergence check fired
+        # immediately, the body never ran, and the *untrained* random init weights
+        # were returned. After the fix the solver must recover the true slope.
+        x, y = _xy(slope=2.0, n=40)
+        algo = offline.LinearRegression(gradient_descent=True, max_iter=20000, learning_rate=1e-3)
+        w = np.asarray(bm.as_jax(algo(y, x))).flatten()
+        assert np.allclose(w[0], 2.0, atol=1e-2), f'did not fit: w={w}'
+
     def test_predict_and_init_weights(self):
         x, y = _xy(slope=2.0)
         algo = offline.LinearRegression()
@@ -173,6 +183,24 @@ def test_call_runs_and_separates(self):
         acc = np.mean((pred.reshape(-1) > 0.5) == y.reshape(-1))
         assert acc >= 0.8
 
+    def test_newton_path_runs_and_fits(self):
+        # M5 regression: the Newton (``gradient_descent=False``) path built a
+        # ``brainpy.Array`` (``bm.zeros`` + in-place diagonal assignment) inside the
+        # ``while_loop`` body, which current JAX rejects with
+        # "Triggering __jax_array__() during abstractification is no longer supported".
+        rng = np.random.RandomState(2)
+        X = rng.randn(200, 3).astype(np.float32)
+        logits = X @ np.array([0.8, -0.5, 0.3], np.float32)
+        p = 1.0 / (1.0 + np.exp(-logits))
+        y = (rng.rand(200) < p).astype(np.float32)[:, None]  # non-separable (finite MLE)
+        algo = offline.LogisticRegression(gradient_descent=False, max_iter=100)
+        w = np.asarray(bm.as_jax(algo(bm.asarray(y), bm.asarray(X)))).reshape(-1)
+        assert w.shape == (3,)
+        assert np.all(np.isfinite(w))
+        # recovers the sign of the true coefficients (0.8, -0.5, 0.3)
+        assert np.sign(w[0]) == 1
+        assert np.sign(w[1]) == -1
+
     def test_predict_applies_sigmoid(self):
         # ``predict`` itself works in isolation (it does not hit the broken call).
         algo = offline.LogisticRegression(max_iter=10)

brainpy/connect/boost_connect_test.py

@@ -58,10 +58,16 @@ def test_fixedprob_include_self_false_and_pre_ratio():
     # no self connections
     assert bool(np.all(np.asarray(pre) != np.asarray(post)))
     indices, indptr = fp.build_csr()
-    # with pre_ratio < 1 only the selected pre rows appear in the indptr
-    # (int(30 * 0.5) = 15 rows -> 16 indptr entries), and counts are monotone.
-    assert indptr.shape[0] == int(30 * 0.5) + 1
+    # A valid CSR indptr always spans the FULL pre range (pre_num + 1 entries);
+    # with pre_ratio < 1 the non-selected pre rows simply have zero out-degree
+    # (equal consecutive indptr values). The previous truncated indptr of length
+    # int(pre_num * pre_ratio) + 1 was a malformed CSR (H4, audit 2026-07-08).
+    assert indptr.shape[0] == 30 + 1
     assert bool(np.all(np.diff(np.asarray(indptr)) >= 0))
+    # indptr is internally consistent: its last entry equals the number of CSR
+    # indices (edges). (``build_coo`` re-samples per call, so we do not cross-check
+    # against the separate ``pre`` draw above.)
+    assert int(np.asarray(indptr)[-1]) == np.asarray(indices).shape[0]
     mat = fp.build_mat()
     # diagonal cleared
     assert not bool(np.any(np.diagonal(np.asarray(mat))))

brainpy/connect/random_conn.py

@@ -133,6 +133,14 @@ def build_coo(self):
         return selected_pre_ids.astype(get_idx_type()), selected_post_ids.astype(get_idx_type())
 
     def build_csr(self):
+        if self.pre_ratio < 1.:
+            # Only a subset of pre neurons is selected, so the CSR index pointer must
+            # still span the FULL pre range (non-selected rows have zero out-degree).
+            # The bespoke ``cumsum`` below produced only ``pre_num_to_select + 1``
+            # entries, giving a structurally-invalid CSR (indptr length !=
+            # pre_num + 1). Route the subset case through ``coo2csr`` over the full
+            # pre range instead (H4, audit 2026-07-08).
+            return coo2csr(self.build_coo(), self.pre_num)
         pre_num_to_select, post_num_to_select, selected_post_ids, pre_ids = self._iii()
         pre_nums = jnp.ones(pre_num_to_select) * post_num_to_select
         if not self.include_self:
@@ -188,14 +196,24 @@ def __init__(self,
 
     def build_coo(self):
         mat_element_num = self.pre_num * self.post_num
-        if self.num > mat_element_num:
+        # A float ``num`` in [0, 1] is documented/allowed by ``__init__`` and denotes
+        # a *fraction* of the all-to-all connection count; coerce it to an absolute
+        # integer count here since the random routines below need an integer size.
+        # Previously a float ``num`` reached ``randint``/``choice`` and raised
+        # ``TypeError: 'float' object cannot be interpreted as an integer``
+        # (H4, audit 2026-07-08).
+        if isinstance(self.num, float):
+            num = int(round(self.num * mat_element_num))
+        else:
+            num = self.num
+        if num > mat_element_num:
             raise ConnectorError(f'"num" must be smaller than "all2all num", '
-                                 f'but got {self.num} > {mat_element_num}')
+                                 f'but got {num} > {mat_element_num}')
         if self.allow_multi_conn:
-            selected_pre_ids = self.rng.randint(0, self.pre_num, (self.num,))
-            selected_post_ids = self.rng.randint(0, self.post_num, (self.num,))
+            selected_pre_ids = self.rng.randint(0, self.pre_num, (num,))
+            selected_post_ids = self.rng.randint(0, self.post_num, (num,))
         else:
-            index = self.rng.choice(mat_element_num, size=(self.num,), replace=False)
+            index = self.rng.choice(mat_element_num, size=(num,), replace=False)
             selected_pre_ids = index // self.post_num
             selected_post_ids = index % self.post_num
         return selected_pre_ids.astype(get_idx_type()), selected_post_ids.astype(get_idx_type())

brainpy/connect/random_conn_coverage_test.py

@@ -92,8 +92,11 @@ def test_fixedprob_include_self_false_pure_python(no_numba):
 def test_fixedprob_pre_ratio_pure_python(no_numba):
     fp = rc.FixedProb(0.5, pre_ratio=0.5, seed=3)((10,), (10,))
     indices, indptr = fp.build_csr()
-    # only int(10 * 0.5) = 5 selected pre rows -> 6 indptr entries
-    assert indptr.shape[0] == int(10 * 0.5) + 1
+    # A valid CSR indptr spans the FULL pre range (pre_num + 1 entries); non-selected
+    # pre rows have zero out-degree. The previous truncated length of
+    # int(pre_num * pre_ratio) + 1 was a malformed CSR (H4, audit 2026-07-08).
+    assert indptr.shape[0] == 10 + 1
+    assert int(np.asarray(indptr)[-1]) == np.asarray(indices).shape[0]
 
 
 def test_fixedprob_allow_multi_conn_jax_path():

brainpy/connect/random_conn_test.py

@@ -40,6 +40,44 @@ def test_require_method(self):
         mat = conn2.require(10, 20, bp.connect.CONN_MAT)
         self.assertTrue(mat.shape == (10, 20))
 
+    def test_build_csr_valid_with_pre_ratio(self):
+        # Regression for H4 (audit 2026-07-08): with ``pre_ratio < 1`` the CSR index
+        # pointer only spanned the *selected* pre neurons (len == pre_num_to_select+1)
+        # instead of the full pre range (len == pre_num+1), producing an invalid CSR.
+        import numpy as np
+        pre_num = 10
+        conn = bp.connect.FixedProb(prob=0.5, pre_ratio=0.5, seed=42, allow_multi_conn=True)
+        conn(pre_size=pre_num, post_size=10)
+        indices, indptr = conn.build_csr()
+        indices = np.asarray(indices)
+        indptr = np.asarray(indptr)
+        self.assertEqual(len(indptr), pre_num + 1)
+        self.assertEqual(int(indptr[0]), 0)
+        self.assertTrue(np.all(np.diff(indptr) >= 0))
+        self.assertEqual(int(indptr[-1]), len(indices))
+        self.assertTrue(np.all((indices >= 0) & (indices < 10)))
+        # Exactly ``int(pre_num * pre_ratio)`` rows carry out-going connections.
+        nonempty_rows = int(np.sum(np.diff(indptr) > 0))
+        self.assertEqual(nonempty_rows, int(pre_num * 0.5))
+
+
+class TestFixedTotalNum(unittest.TestCase):
+    def test_float_num_is_fraction_of_all2all(self):
+        # Regression for H4 (audit 2026-07-08): a float ``num`` (documented as allowed
+        # in [0, 1]) reached ``randint``/``choice`` and raised ``TypeError`` because a
+        # float cannot be a shape. It now denotes a fraction of the all-to-all count.
+        conn = bp.connect.FixedTotalNum(num=0.1, seed=1)
+        conn(pre_size=10, post_size=10)  # all2all == 100
+        pre_ids, post_ids = conn.build_coo()
+        self.assertEqual(len(pre_ids), 10)
+        self.assertEqual(len(post_ids), 10)
+
+    def test_int_num_unchanged(self):
+        conn = bp.connect.FixedTotalNum(num=25, seed=1)
+        conn(pre_size=10, post_size=10)
+        pre_ids, _ = conn.build_coo()
+        self.assertEqual(len(pre_ids), 25)
+
 
 def test_random_fix_pre1():
     for num in [0.4, 20]:

brainpy/connect/regular_conn.py

@@ -192,6 +192,18 @@ def f_connect(pre_id):
             strides = jnp.asarray(get_size_length(self.post_size))
             pres = jnp.sum(pres * strides, axis=1)
             posts = jnp.sum(posts * strides, axis=1)
+        if self.periodic_boundary:
+            # On small grids the periodic wrap (``post_ids % sizes``) can map several
+            # strides onto the same post neuron, producing duplicate ``(pre, post)``
+            # edges (e.g. a 2x2 GridFour yields each edge twice). De-duplicate while
+            # preserving order (M1, audit 2026-07-08).
+            pres_np = np.asarray(pres)
+            posts_np = np.asarray(posts)
+            _, uniq_idx = np.unique(np.stack([pres_np, posts_np], axis=1),
+                                    axis=0, return_index=True)
+            uniq_idx = np.sort(uniq_idx)
+            pres = pres_np[uniq_idx]
+            posts = posts_np[uniq_idx]
         return jnp.asarray(pres, dtype=get_idx_type()), jnp.asarray(posts, dtype=get_idx_type())
 
 

brainpy/connect/regular_conn_test.py

@@ -116,3 +116,21 @@ def test_grid_N(self):
 
                     print(f'periodic_boundary = {periodic_boundary}, include_self = {include_self}, size = {size}')
                     self.assertTrue(bp.math.allclose(mat, new_mat))
+
+    def test_periodic_boundary_no_duplicate_edges(self):
+        # Regression for M1 (audit 2026-07-08): on small grids the periodic wrap
+        # (``post_ids % sizes``) maps several strides onto the same post neuron,
+        # producing duplicate (pre, post) COO edges (e.g. a 2x2 GridFour yields every
+        # edge twice). Duplicates are invisible in a boolean matrix but double-count
+        # weights in COO-based synapses.
+        for cls, kwargs in [(bp.conn.GridFour, {}),
+                            (bp.conn.GridEight, {}),
+                            (bp.conn.GridN, {'N': 2})]:
+            for size in [(2, 2), (2, 3), (3, 3)]:
+                conn = cls(periodic_boundary=True, **kwargs)(size, size)
+                pre_ids, post_ids = conn.build_coo()
+                pairs = np.stack([np.asarray(pre_ids), np.asarray(post_ids)], axis=1)
+                n_unique = len(np.unique(pairs, axis=0))
+                self.assertEqual(len(pairs), n_unique,
+                                 msg=f'{cls.__name__} size={size}: '
+                                     f'{len(pairs) - n_unique} duplicate edges')

brainpy/dnn/normalization.py

@@ -557,7 +557,9 @@ def update(self, x):
         if x.shape[-len(self.normalized_shape):] != self.normalized_shape:
             raise ValueError(f'Expect the input shape should be (..., {", ".join(map(str, self.normalized_shape))}), '
                              f'but we got {x.shape}')
-        axis = tuple(range(0, x.ndim - len(self.normalized_shape)))
+        # Normalize over the trailing ``normalized_shape`` dimensions (LayerNorm's
+        # definition), NOT the leading batch/time axes (C1, audit 2026-07-08).
+        axis = tuple(range(x.ndim - len(self.normalized_shape), x.ndim))
         mean = jnp.mean(bm.as_jax(x), axis=axis, keepdims=True)
         variance = jnp.var(bm.as_jax(x), axis=axis, keepdims=True)
         inv = lax.rsqrt(variance + lax.convert_element_type(self.epsilon, x.dtype))

brainpy/dnn/normalization_test.py

@@ -118,6 +118,30 @@ def test_BatchNorm_batch_is_biased_normalized(self):
         self.assertTrue(bool(jnp.allclose(out.mean(axis=(0, 1)), 0.0, atol=1e-5)))
         self.assertTrue(bool(jnp.allclose(out.var(axis=(0, 1)), 1.0, atol=1e-4)))
 
+    def test_LayerNorm_normalizes_over_trailing_axes(self):
+        # Regression for C1 (audit 2026-07-08): ``LayerNorm`` must normalize over the
+        # trailing ``normalized_shape`` dimensions, not the leading (batch/time) axes.
+        # With the bug it reduced over ``range(0, ndim - len(normalized_shape))`` and
+        # produced results ~1.7 away from the reference.
+        import jax.numpy as jnp
+        bm.random.seed(2718)
+        normalized_shape = (5, 10)
+        net = bp.dnn.LayerNorm(normalized_shape, elementwise_affine=False, mode=bm.training_mode)
+        x = bm.as_jax(bm.random.randn(4, 3, 5, 10) * 3.0 + 2.0)
+        out = bm.as_jax(net(x))
+
+        # Reference: normalize each (5, 10) block independently over its last 2 dims.
+        axes = (-2, -1)
+        mean = jnp.mean(x, axis=axes, keepdims=True)
+        var = jnp.var(x, axis=axes, keepdims=True)
+        ref = (x - mean) / jnp.sqrt(var + 1e-5)
+        self.assertTrue(bool(jnp.allclose(out, ref, atol=1e-4)),
+                        msg=f'max|out-ref|={float(jnp.max(jnp.abs(out - ref)))}')
+
+        # Each normalized block must be (approximately) zero-mean / unit-variance.
+        self.assertTrue(bool(jnp.allclose(jnp.mean(out, axis=axes), 0.0, atol=1e-4)))
+        self.assertTrue(bool(jnp.allclose(jnp.var(out, axis=axes), 1.0, atol=1e-3)))
+
 
 if __name__ == '__main__':
     absltest.main()