fix: prevent MilsteinGradFree OOM crash; make GPU test suite hardware-independent (#837)

* fix: prevent MilsteinGradFree OOM crash; make GPU test suite hardware-independent

The Linux CI job (`pytest brainpy/`) was OOM-killed (exit 143) at
`brainpy/integrators/sde/normal_coverage_test.py` [69%]. Root cause: the
vector-Wiener branch of `MilsteinGradFree.step` used the full `(m, m)`
diffusion-bar matrix together with `minus * jnp.sum(noise_p2, -1)`, which left
the noise dimension in the integrated state. Each step grew the output by two
axes — `()` -> `(m, m)` -> `(m, m, m, m)` -> ... — a multi-GB blow-up over a long
integration that exhausted the runner's RAM (macOS/Windows runners had enough
headroom to survive, so only Linux crashed). Take the diagonal of the
diffusion-bar block and contract the per-component Milstein correction over the
noise axis, mirroring the shape-preserving pattern already used by `Milstein`.
The integrated state now stays correctly shaped and the full suite peaks at a
bounded ~9 GB with the crash region flat (no spike).

Also make the suite pass on GPU developer machines without changing CPU/CI
behaviour:

* conftest.py: pin `jax_default_matmul_precision='highest'`. On NVIDIA GPUs the
  default uses TF32 for float32 matmuls (~1e-4 relative error), which broke the
  operator-vs-dense correctness comparisons (JIT-connectivity layers,
  orthonormality checks) on GPU while passing on CPU.
* brainpy/dnn/linear_test.py: use float32-appropriate tolerances
  (`rtol=1e-4, atol=1e-5`) for the JIT-operator vs dense `x @ conn` comparison;
  the default `atol=1e-8` is tighter than float32 rounding for the near-zero
  symmetric-uniform outputs.
* brainpy/math/object_transform/object_transform_fixes_test.py: guard the
  `.cuda()` / `.tpu()` `RuntimeError` assertions on device availability so the
  test stays meaningful on CPU-only CI yet does not fail on GPU/TPU machines.

* fix(deps): require braintools>=0.3.0 (0.2.0 was yanked from PyPI)

braintools 0.2.0 was yanked from PyPI, so the previous `braintools>=0.2.0`
pin became unsatisfiable and broke `pip install -r requirements.txt` on every
CI runner (install failed in ~10s, before any test ran). 0.3.0 is the next
released version and carries the surrogate / metric fixes that the
`brainpy.math.surrogate` and L1-loss tests assert. Bump the pin in both
requirements.txt and pyproject.toml.

* refactor(math): reuse brainunit.math einops; drop the local port

brainpy's ein_reduce / ein_rearrange / ein_repeat / ein_shape duplicated the
einops implementation that now lives in brainunit.math (einreduce / einrearrange
/ einrepeat / einshape — behaviour-identical and accepting brainpy ``Array``
instances directly). Re-export the historical ``ein_*`` names as thin aliases of
brainunit's and delete the duplicated implementation (einops.py,
einops_parsing.py) and its dedicated tests (einops_test.py,
einops_coverage_test.py, einops_parsing_test.py). The einops tests in
math_compat_fixes_test.py now exercise the public ``bm.ein_*`` aliases and assert
the re-export wiring.

* chore(math): remove taichi/tifunc remnants

The taichi backend is gone; drop the fully-skipped tifunc_test.py (a taichi
``ti.kernel`` test) and the stale ``method`` / Taichi paragraph in the csrmv
docstring (that parameter no longer exists — csrmv dispatches through brainevent).

Author: chaoming0625

brainpy/dnn/linear_test.py

@@ -150,7 +150,10 @@ def test_JitFPHomoLinear(self, prob, weight, shape):
         self.assertTrue(y.shape == shape + (200,))
 
         conn_matrix = f.get_conn_matrix()
-        self.assertTrue(bm.allclose(y, x @ conn_matrix.T))
+        # float32-appropriate tolerances: the JIT operator and the dense ``x @ conn``
+        # differ at float32 rounding level; the default ``atol=1e-8`` is tighter than that
+        # for the symmetric-uniform layer whose outputs sit near zero.
+        self.assertTrue(bm.allclose(y, x @ conn_matrix.T, rtol=1e-4, atol=1e-5))
         # print(conn_matrix.shape)
         # self.assertTrue(conn_matrix.shape == (200, 100))
 
@@ -168,7 +171,10 @@ def test_JitFPUniformLinear(self, prob, w_low, w_high, shape):
         self.assertTrue(y.shape == shape + (200,))
 
         conn_matrix = f.get_conn_matrix()
-        self.assertTrue(bm.allclose(y, x @ conn_matrix.T))
+        # float32-appropriate tolerances: the JIT operator and the dense ``x @ conn``
+        # differ at float32 rounding level; the default ``atol=1e-8`` is tighter than that
+        # for the symmetric-uniform layer whose outputs sit near zero.
+        self.assertTrue(bm.allclose(y, x @ conn_matrix.T, rtol=1e-4, atol=1e-5))
 
     @parameterized.product(
         prob=[0.1],
@@ -184,7 +190,10 @@ def test_JitFPNormalLinear(self, prob, w_mu, w_sigma, shape):
         self.assertTrue(y.shape == shape + (200,))
 
         conn_matrix = f.get_conn_matrix()
-        self.assertTrue(bm.allclose(y, x @ conn_matrix.T))
+        # float32-appropriate tolerances: the JIT operator and the dense ``x @ conn``
+        # differ at float32 rounding level; the default ``atol=1e-8`` is tighter than that
+        # for the symmetric-uniform layer whose outputs sit near zero.
+        self.assertTrue(bm.allclose(y, x @ conn_matrix.T, rtol=1e-4, atol=1e-5))
 
     @parameterized.product(
         prob=[0.1],
@@ -202,7 +211,10 @@ def test_EventJitFPHomoLinear(self, prob, weight, shape):
         self.assertTrue(y2.shape == shape + (200,))
 
         conn_matrix = f.get_conn_matrix()
-        self.assertTrue(bm.allclose(y, x @ conn_matrix.T))
+        # float32-appropriate tolerances: the JIT operator and the dense ``x @ conn``
+        # differ at float32 rounding level; the default ``atol=1e-8`` is tighter than that
+        # for the symmetric-uniform layer whose outputs sit near zero.
+        self.assertTrue(bm.allclose(y, x @ conn_matrix.T, rtol=1e-4, atol=1e-5))
 
     @parameterized.product(
         prob=[0.1],
@@ -221,7 +233,10 @@ def test_EventJitFPUniformLinear(self, prob, w_low, w_high, shape):
         self.assertTrue(y2.shape == shape + (200,))
 
         conn_matrix = f.get_conn_matrix()
-        self.assertTrue(bm.allclose(y, x @ conn_matrix.T))
+        # float32-appropriate tolerances: the JIT operator and the dense ``x @ conn``
+        # differ at float32 rounding level; the default ``atol=1e-8`` is tighter than that
+        # for the symmetric-uniform layer whose outputs sit near zero.
+        self.assertTrue(bm.allclose(y, x @ conn_matrix.T, rtol=1e-4, atol=1e-5))
 
     @parameterized.product(
         prob=[0.1],
@@ -240,7 +255,10 @@ def test_EventJitFPNormalLinear(self, prob, w_mu, w_sigma, shape):
         self.assertTrue(y2.shape == shape + (200,))
 
         conn_matrix = f.get_conn_matrix()
-        self.assertTrue(bm.allclose(y, x @ conn_matrix.T))
+        # float32-appropriate tolerances: the JIT operator and the dense ``x @ conn``
+        # differ at float32 rounding level; the default ``atol=1e-8`` is tighter than that
+        # for the symmetric-uniform layer whose outputs sit near zero.
+        self.assertTrue(bm.allclose(y, x @ conn_matrix.T, rtol=1e-4, atol=1e-5))
 
 
 if __name__ == '__main__':

brainpy/integrators/sde/normal.py

@@ -502,10 +502,20 @@ def step(self, *args, **kwargs):
                 else:
                     integral += diffusions[key] * noise
                 noise_p2 = (noise ** 2 - dt) if self.intg_type == constants.ITO_SDE else noise ** 2
-                minus = (diffusion_bars[key] - diffusions[key]) / 2 / jnp.sqrt(dt)
                 if self.wiener_type == constants.VECTOR_WIENER:
-                    integral += minus * jnp.sum(noise_p2, axis=-1)
+                    # ``y_bars[key]`` carries the noise axis (one support value per noise
+                    # component ``j``: ``y_bar_j = Y + f dt + g_j sqrt(dt)``), so
+                    # ``diffusion_bars[key]`` has a trailing ``(m, m)`` block whose diagonal is
+                    # ``g_j(y_bar_j)``. Previously the full ``(m, m)`` matrix was used together
+                    # with ``minus * jnp.sum(noise_p2, -1)``, which left the noise dimension in
+                    # the state and grew the output by two axes every step (a multi-GB blow-up
+                    # for long integrations). Take the diagonal and contract the per-component
+                    # Milstein correction over the noise axis, mirroring ``Milstein``.
+                    g_bar = jnp.diagonal(diffusion_bars[key], axis1=-2, axis2=-1)
+                    minus = (g_bar - diffusions[key]) / 2 / jnp.sqrt(dt)
+                    integral += jnp.sum(minus * noise_p2, axis=-1)
                 else:
+                    minus = (diffusion_bars[key] - diffusions[key]) / 2 / jnp.sqrt(dt)
                     integral += minus * noise_p2
             integrals.append(integral)
         return integrals if len(self.variables) > 1 else integrals[0]

brainpy/math/__init__.py

@@ -62,7 +62,14 @@
 from .compat_tensorflow import *
 from .datatypes import *
 from .delayvars import *
-from .einops import *
+# einops-style helpers are reused from ``brainunit.math`` (the local port was
+# removed); keep the historical ``ein_*`` names as thin aliases.
+from brainunit.math import (
+    einreduce as ein_reduce,
+    einrearrange as ein_rearrange,
+    einrepeat as ein_repeat,
+    einshape as ein_shape,
+)
 from .environment import *
 from .interoperability import *
 # environment settings