brainstate 0.0.2.post20240913__py2.py3-none-any.whl → 0.0.2.post20241009__py2.py3-none-any.whl

brainstate/functional/_activations_test.py

@@ -0,0 +1,331 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Tests for nn module."""
+
+import itertools
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+import scipy.stats
+from absl.testing import parameterized
+from jax._src import test_util as jtu
+from jax.test_util import check_grads
+
+import brainstate as bst
+
+
+class NNFunctionsTest(jtu.JaxTestCase):
+  @jtu.skip_on_flag("jax_skip_slow_tests", True)
+  def testSoftplusGrad(self):
+    check_grads(bst.functional.softplus, (1e-8,), order=4, )
+
+  def testSoftplusGradZero(self):
+    check_grads(bst.functional.softplus, (0.,), order=1)
+
+  def testSoftplusGradInf(self):
+    self.assertAllClose(1., jax.grad(bst.functional.softplus)(float('inf')))
+
+  def testSoftplusGradNegInf(self):
+    check_grads(bst.functional.softplus, (-float('inf'),), order=1)
+
+  def testSoftplusGradNan(self):
+    check_grads(bst.functional.softplus, (float('nan'),), order=1)
+
+  @parameterized.parameters([int, float] + jtu.dtypes.floating + jtu.dtypes.integer)
+  def testSoftplusZero(self, dtype):
+    self.assertEqual(jnp.log(dtype(2)), bst.functional.softplus(dtype(0)))
+
+  def testSparseplusGradZero(self):
+    check_grads(bst.functional.sparse_plus, (-2.,), order=1)
+
+  def testSparseplusGrad(self):
+    check_grads(bst.functional.sparse_plus, (0.,), order=1)
+
+  def testSparseplusAndSparseSigmoid(self):
+    self.assertAllClose(
+      jax.grad(bst.functional.sparse_plus)(0.),
+      bst.functional.sparse_sigmoid(0.),
+      check_dtypes=False)
+    self.assertAllClose(
+      jax.grad(bst.functional.sparse_plus)(2.),
+      bst.functional.sparse_sigmoid(2.),
+      check_dtypes=False)
+    self.assertAllClose(
+      jax.grad(bst.functional.sparse_plus)(-2.),
+      bst.functional.sparse_sigmoid(-2.),
+      check_dtypes=False)
+
+  def testSquareplusGrad(self):
+    check_grads(bst.functional.squareplus, (1e-8,), order=4,
+                )
+
+  def testSquareplusGradZero(self):
+    check_grads(bst.functional.squareplus, (0.,), order=1,
+                )
+
+  def testSquareplusGradNegInf(self):
+    check_grads(bst.functional.squareplus, (-float('inf'),), order=1,
+                )
+
+  def testSquareplusGradNan(self):
+    check_grads(bst.functional.squareplus, (float('nan'),), order=1,
+                )
+
+  @parameterized.parameters([float] + jtu.dtypes.floating)
+  def testSquareplusZero(self, dtype):
+    self.assertEqual(dtype(1), bst.functional.squareplus(dtype(0), dtype(4)))
+
+  def testMishGrad(self):
+    check_grads(bst.functional.mish, (1e-8,), order=4,
+                )
+
+  def testMishGradZero(self):
+    check_grads(bst.functional.mish, (0.,), order=1,
+                )
+
+  def testMishGradNegInf(self):
+    check_grads(bst.functional.mish, (-float('inf'),), order=1,
+                )
+
+  def testMishGradNan(self):
+    check_grads(bst.functional.mish, (float('nan'),), order=1,
+                )
+
+  @parameterized.parameters([float] + jtu.dtypes.floating)
+  def testMishZero(self, dtype):
+    self.assertEqual(dtype(0), bst.functional.mish(dtype(0)))
+
+  def testReluGrad(self):
+    rtol = None
+    check_grads(bst.functional.relu, (1.,), order=3, rtol=rtol)
+    check_grads(bst.functional.relu, (-1.,), order=3, rtol=rtol)
+    jaxpr = jax.make_jaxpr(jax.grad(bst.functional.relu))(0.)
+    self.assertGreaterEqual(len(jaxpr.jaxpr.eqns), 2)
+
+  def testRelu6Grad(self):
+    rtol = None
+    check_grads(bst.functional.relu6, (1.,), order=3, rtol=rtol)
+    check_grads(bst.functional.relu6, (-1.,), order=3, rtol=rtol)
+    self.assertAllClose(jax.grad(bst.functional.relu6)(0.), 0., check_dtypes=False)
+    self.assertAllClose(jax.grad(bst.functional.relu6)(6.), 0., check_dtypes=False)
+
+  def testSoftplusValue(self):
+    val = bst.functional.softplus(89.)
+    self.assertAllClose(val, 89., check_dtypes=False)
+
+  def testSparseplusValue(self):
+    val = bst.functional.sparse_plus(89.)
+    self.assertAllClose(val, 89., check_dtypes=False)
+
+  def testSparsesigmoidValue(self):
+    self.assertAllClose(bst.functional.sparse_sigmoid(-2.), 0., check_dtypes=False)
+    self.assertAllClose(bst.functional.sparse_sigmoid(2.), 1., check_dtypes=False)
+    self.assertAllClose(bst.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
+
+  def testSquareplusValue(self):
+    val = bst.functional.squareplus(1e3)
+    self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
+
+  def testMishValue(self):
+    val = bst.functional.mish(1e3)
+    self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
+
+  def testEluValue(self):
+    val = bst.functional.elu(1e4)
+    self.assertAllClose(val, 1e4, check_dtypes=False)
+
+  def testGluValue(self):
+    val = bst.functional.glu(jnp.array([1.0, 0.0]), axis=0)
+    self.assertAllClose(val, jnp.array([0.5]))
+
+  @parameterized.parameters(False, True)
+  def testGeluIntType(self, approximate):
+    val_float = bst.functional.gelu(jnp.array(-1.0), approximate=approximate)
+    val_int = bst.functional.gelu(jnp.array(-1), approximate=approximate)
+    self.assertAllClose(val_float, val_int)
+
+  @parameterized.parameters(False, True)
+  def testGelu(self, approximate):
+    def gelu_reference(x):
+      return x * scipy.stats.norm.cdf(x)
+
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [rng((4, 5, 6), jnp.float32)]
+    self._CheckAgainstNumpy(
+      gelu_reference, partial(bst.functional.gelu, approximate=approximate), args_maker,
+      check_dtypes=False, tol=1e-3 if approximate else None)
+
+  @parameterized.parameters(*itertools.product(
+    (jnp.float32, jnp.bfloat16, jnp.float16),
+    (partial(bst.functional.gelu, approximate=False),
+     partial(bst.functional.gelu, approximate=True),
+     bst.functional.relu,
+     bst.functional.softplus,
+     bst.functional.sparse_plus,
+     bst.functional.sigmoid,
+     bst.functional.squareplus,
+     bst.functional.mish)))
+  def testDtypeMatchesInput(self, dtype, fn):
+    x = jnp.zeros((), dtype=dtype)
+    out = fn(x)
+    self.assertEqual(out.dtype, dtype)
+
+  def testEluMemory(self):
+    # see https://github.com/google/jax/pull/1640
+    with jax.enable_checks(False):  # With checks we materialize the array
+      jax.make_jaxpr(lambda: bst.functional.elu(jnp.ones((10 ** 12,))))  # don't oom
+
+  def testHardTanhMemory(self):
+    # see https://github.com/google/jax/pull/1640
+    with jax.enable_checks(False):  # With checks we materialize the array
+      jax.make_jaxpr(lambda: bst.functional.hard_tanh(jnp.ones((10 ** 12,))))  # don't oom
+
+  @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+  def testSoftmaxEmptyArray(self, fn):
+    x = jnp.array([], dtype=float)
+    self.assertArraysEqual(fn(x), x)
+
+  @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+  def testSoftmaxEmptyMask(self, fn):
+    x = jnp.array([5.5, 1.3, -4.2, 0.9])
+    m = jnp.zeros_like(x, dtype=bool)
+    expected = jnp.full_like(x, 0.0 if fn is bst.functional.softmax else -jnp.inf)
+    self.assertArraysEqual(fn(x, where=m), expected)
+
+  @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+  def testSoftmaxWhereMask(self, fn):
+    x = jnp.array([5.5, 1.3, -4.2, 0.9])
+    m = jnp.array([True, False, True, True])
+
+    out = fn(x, where=m)
+    self.assertAllClose(out[m], fn(x[m]))
+
+    probs = out if fn is bst.functional.softmax else jnp.exp(out)
+    self.assertAllClose(probs.sum(), 1.0)
+
+  @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+  def testSoftmaxWhereGrad(self, fn):
+    # regression test for https://github.com/google/jax/issues/19490
+    x = jnp.array([36., 10000.])
+    mask = x < 1000
+
+    f = lambda x, mask: fn(x, where=mask)[0]
+
+    self.assertAllClose(jax.grad(f)(x, mask), jnp.zeros_like(x))
+
+  def testSoftmaxGrad(self):
+    x = jnp.array([5.5, 1.3, -4.2, 0.9])
+    jtu.check_grads(bst.functional.softmax, (x,), order=2, atol=5e-3)
+
+  def testStandardizeWhereMask(self):
+    x = jnp.array([5.5, 1.3, -4.2, 0.9])
+    m = jnp.array([True, False, True, True])
+    x_filtered = jnp.take(x, jnp.array([0, 2, 3]))
+
+    out_masked = jnp.take(bst.functional.standardize(x, where=m), jnp.array([0, 2, 3]))
+    out_filtered = bst.functional.standardize(x_filtered)
+
+    self.assertAllClose(out_masked, out_filtered)
+
+  def testOneHot(self):
+    actual = bst.functional.one_hot(jnp.array([0, 1, 2]), 3)
+    expected = jnp.array([[1., 0., 0.],
+                          [0., 1., 0.],
+                          [0., 0., 1.]])
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+    actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3)
+    expected = jnp.array([[0., 1., 0.],
+                          [0., 0., 1.],
+                          [1., 0., 0.]])
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+  def testOneHotOutOfBound(self):
+    actual = bst.functional.one_hot(jnp.array([-1, 3]), 3)
+    expected = jnp.array([[0., 0., 0.],
+                          [0., 0., 0.]])
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+  def testOneHotNonArrayInput(self):
+    actual = bst.functional.one_hot([0, 1, 2], 3)
+    expected = jnp.array([[1., 0., 0.],
+                          [0., 1., 0.],
+                          [0., 0., 1.]])
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+  def testOneHotCustomDtype(self):
+    actual = bst.functional.one_hot(jnp.array([0, 1, 2]), 3, dtype=jnp.bool_)
+    expected = jnp.array([[True, False, False],
+                          [False, True, False],
+                          [False, False, True]])
+    self.assertAllClose(actual, expected)
+
+  def testOneHotAxis(self):
+    expected = jnp.array([[0., 1., 0.],
+                          [0., 0., 1.],
+                          [1., 0., 0.]]).T
+
+    actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=0)
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+    actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=-2)
+    self.assertAllClose(actual, expected, check_dtypes=False)
+
+  def testTanhExists(self):
+    print(bst.functional.tanh)  # doesn't crash
+
+  def testCustomJVPLeak(self):
+    # https://github.com/google/jax/issues/8171
+    @jax.jit
+    def fwd():
+      a = jnp.array(1.)
+
+      def f(hx, _):
+        hx = bst.functional.sigmoid(hx + a)
+        return hx, None
+
+      hx = jnp.array(0.)
+      jax.lax.scan(f, hx, None, length=2)
+
+    with jax.checking_leaks():
+      fwd()  # doesn't crash
+
+  def testCustomJVPLeak2(self):
+    # https://github.com/google/jax/issues/8171
+    # The above test uses jax.bst.functional.sigmoid, as in the original #8171, but that
+    # function no longer actually has a custom_jvp! So we inline the old def.
+
+    @jax.custom_jvp
+    def sigmoid(x):
+      one = jnp.float32(1)
+      return jax.lax.div(one, jax.lax.add(one, jax.lax.exp(jax.lax.neg(x))))
+
+    sigmoid.defjvps(lambda g, ans, x: g * ans * (jnp.float32(1) - ans))
+
+    @jax.jit
+    def fwd():
+      a = jnp.array(1., 'float32')
+
+      def f(hx, _):
+        hx = sigmoid(hx + a)
+        return hx, None
+
+      hx = jnp.array(0., 'float32')
+      jax.lax.scan(f, hx, None, length=2)
+
+    with jax.checking_leaks():
+      fwd()  # doesn't crash

brainstate/functional/_normalization.py

@@ -15,12 +15,13 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Union
 
+import brainunit as u
 import jax
-import jax.numpy as jnp
 
-from .._utils import set_module_as
+from brainstate._utils import set_module_as
+from brainstate.typing import ArrayLike
 
 __all__ = [
   'weight_standardization',
@@ -29,18 +30,18 @@ __all__ = [
 
 @set_module_as('brainstate.functional')
 def weight_standardization(
-    w: jax.typing.ArrayLike,
+    w: ArrayLike,
     eps: float = 1e-4,
     gain: Optional[jax.Array] = None,
     out_axis: int = -1,
-):
+) -> Union[jax.Array, u.Quantity]:
   """
   Scaled Weight Standardization,
   see `Micro-Batch Training with Batch-Channel Normalization and Weight Standardization <https://paperswithcode.com/paper/weight-standardization>`_.
 
   Parameters
   ----------
-  w : jax.typing.ArrayLike
+  w : ArrayLike
       The weight tensor.
   eps : float
       A small value to avoid division by zero.
@@ -51,7 +52,7 @@ def weight_standardization(
 
   Returns
   -------
-  jax.typing.ArrayLike
+  ArrayLike
       The scaled weight tensor.
   """
   if out_axis < 0:
@@ -63,9 +64,19 @@ def weight_standardization(
       fan_in *= w.shape[i]
       axes.append(i)
   # normalize the weight
-  mean = jnp.mean(w, axis=axes, keepdims=True)
-  var = jnp.var(w, axis=axes, keepdims=True)
-  scale = jax.lax.rsqrt(jnp.maximum(var * fan_in, eps))
+  mean = u.math.mean(w, axis=axes, keepdims=True)
+  var = u.math.var(w, axis=axes, keepdims=True)
+
+  temp = u.math.maximum(var * fan_in, eps)
+  if isinstance(temp, u.Quantity):
+    unit = temp.unit
+    temp = temp.mantissa
+    if unit.is_unitless:
+      scale = jax.lax.rsqrt(temp)
+    else:
+      scale = u.Quantity(jax.lax.rsqrt(temp), unit=1 / unit ** 0.5)
+  else:
+    scale = jax.lax.rsqrt(temp)
   if gain is not None:
     scale = gain * scale
   shift = mean * scale

brainstate/init/_generic.py

@@ -24,6 +24,7 @@ import numpy as np
 from brainstate._state import State
 from brainstate.typing import ArrayLike
 from ._base import to_size
+from brainstate.mixin import Mode
 
 __all__ = [
   'param',
@@ -36,7 +37,7 @@ def _is_scalar(x):
   return bu.math.isscalar(x)
 
 
-def are_shapes_broadcastable(shape1, shape2):
+def are_broadcastable_shapes(shape1, shape2):
   """
   Check if two shapes are broadcastable.
 
@@ -88,6 +89,7 @@ def param(
     batch_size: Optional[int] = None,
     allow_none: bool = True,
     allow_scalar: bool = True,
+    mode: Mode = None,
 ):
   """Initialize parameters.
 
@@ -143,7 +145,7 @@ def param(
     raise ValueError(f'Unknown parameter type: {type(parameter)}')
 
   # Check if the shape of the parameter matches the given size
-  if not are_shapes_broadcastable(parameter.shape, sizes):
+  if not are_broadcastable_shapes(parameter.shape, sizes):
     raise ValueError(f'The shape of the parameter {parameter.shape} does not match with the given size {sizes}')
 
   # Expand the parameter to match the given batch size

brainstate/mixin.py

@@ -18,7 +18,7 @@
 from typing import Sequence, Optional, TypeVar
 from typing import (_SpecialForm, _type_check, _remove_dups_flatten, _UnionGenericAlias)
 
-from .typing import PyTree
+from brainstate.typing import PyTree
 
 T = TypeVar('T')
 State = None

brainstate/nn/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+from . import metrics
 from ._base import *
 from ._base import __all__ as base_all
 from ._connections import *
@@ -39,6 +40,8 @@ from ._readout import *
 from ._readout import __all__ as readout_all
 from ._synouts import *
 from ._synouts import __all__ as synouts_all
+from .event import *
+from .event import __all__ as event_all
 
 __all__ = (
     base_all +
@@ -53,7 +56,8 @@ __all__ = (
     readout_all +
     synouts_all +
     _projection_all +
-    _misc_all
+    _misc_all +
+    event_all
 )
 
 del (
@@ -68,5 +72,6 @@ del (
   readout_all,
   synouts_all,
   _projection_all,
-  _misc_all
+  _misc_all,
+  event_all
 )

brainstate/nn/_base.py

@@ -20,8 +20,8 @@ from __future__ import annotations
 import inspect
 from typing import Sequence, Optional, Tuple, Union
 
-from .._module import Module, UpdateReturn, Container, visible_module_dict
-from ..mixin import Mixin, DelayedInitializer, DelayedInit
+from brainstate._module import Module, UpdateReturn, Container, visible_module_dict
+from brainstate.mixin import Mixin, DelayedInitializer, DelayedInit
 
 __all__ = [
   'ExplicitInOutSize',

brainstate/nn/_connections.py

@@ -25,10 +25,10 @@ import jax
 import jax.numpy as jnp
 
 from ._base import DnnLayer
-from .. import init, functional
-from .._state import ParamState
-from ..mixin import Mode
-from ..typing import ArrayLike
+from brainstate import init, functional
+from brainstate._state import ParamState
+from brainstate.mixin import Mode
+from brainstate.typing import ArrayLike
 
 T = TypeVar('T')
 

brainstate/nn/_dynamics.py

@@ -24,11 +24,11 @@ import jax.numpy as jnp
 
 from ._base import ExplicitInOutSize
 from ._misc import exp_euler_step
-from .. import environ, init, surrogate
-from .._module import Dynamics
-from .._state import ShortTermState
-from ..mixin import DelayedInit, Mode, AlignPost
-from ..typing import DTypeLike, ArrayLike, Size
+from brainstate import environ, init, surrogate
+from brainstate._module import Dynamics
+from brainstate._state import ShortTermState
+from brainstate.mixin import DelayedInit, Mode, AlignPost
+from brainstate.typing import DTypeLike, ArrayLike, Size
 
 __all__ = [
   # neuron models

brainstate/nn/_elementwise.py

@@ -19,16 +19,16 @@ from __future__ import annotations
 
 from typing import Optional
 
-import brainunit as bu
+import brainunit as u
 import jax.numpy as jnp
 import jax.typing
 
 from ._base import ElementWiseBlock
-from .. import environ, random, functional as F
-from .._module import Module
-from .._state import ParamState
-from ..mixin import Mode
-from ..typing import ArrayLike
+from brainstate import environ, random, functional as F
+from brainstate._module import Module
+from brainstate._state import ParamState
+from brainstate.mixin import Mode
+from brainstate.typing import ArrayLike
 
 __all__ = [
   # activation functions
@@ -83,7 +83,7 @@ class Threshold(Module, ElementWiseBlock):
     self.value = value
 
   def __call__(self, x: ArrayLike) -> ArrayLike:
-    dtype = bu.math.get_dtype(x)
+    dtype = u.math.get_dtype(x)
     return jnp.where(x > jnp.asarray(self.threshold, dtype=dtype),
                      x,
                      jnp.asarray(self.value, dtype=dtype))
@@ -1143,7 +1143,7 @@ class Dropout(Module, ElementWiseBlock):
     self.prob = prob
 
   def __call__(self, x):
-    dtype = bu.math.get_dtype(x)
+    dtype = u.math.get_dtype(x)
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
     if fit_phase and self.prob < 1.:
       keep_mask = random.bernoulli(self.prob, x.shape)
@@ -1173,7 +1173,7 @@ class _DropoutNd(Module, ElementWiseBlock):
     self.channel_axis = channel_axis
 
   def __call__(self, x):
-    dtype = bu.math.get_dtype(x)
+    dtype = u.math.get_dtype(x)
     # get fit phase
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
 

brainstate/nn/_embedding.py

@@ -17,9 +17,9 @@ from typing import Optional, Callable, Union
 
 from ._base import DnnLayer
 from .. import init
-from .._state import ParamState
-from ..mixin import Mode, Training
-from ..typing import ArrayLike
+from brainstate._state import ParamState
+from brainstate.mixin import Mode, Training
+from brainstate.typing import ArrayLike
 
 __all__ = [
   'Embedding',

brainstate/nn/_normalizations.py

@@ -25,9 +25,9 @@ import jax.numpy as jnp
 
 from ._base import DnnLayer
 from .. import environ, init
-from .._state import LongTermState, ParamState
-from ..mixin import Mode
-from ..typing import DTypeLike, ArrayLike, Size, Axes
+from brainstate._state import LongTermState, ParamState
+from brainstate.mixin import Mode
+from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
 
 __all__ = [
   'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d',

brainstate/nn/_others.py

@@ -23,8 +23,8 @@ import brainunit as bu
 import jax.numpy as jnp
 
 from ._base import DnnLayer
-from .. import random, environ, typing, init
-from ..mixin import Mode
+from brainstate.mixin import Mode
+from brainstate import random, environ, typing, init
 
 __all__ = [
   'DropoutFixed',

brainstate/nn/_poolings.py

@@ -21,7 +21,7 @@ import functools
 from typing import Sequence, Optional
 from typing import Union, Tuple, Callable, List
 
-import brainunit as bu
+import brainunit as u
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -29,7 +29,7 @@ import numpy as np
 from ._base import DnnLayer, ExplicitInOutSize
 from .. import environ
 from ..mixin import Mode
-from ..typing import Size
+from brainstate.typing import Size
 
 __all__ = [
   'Flatten', 'Unflatten',
@@ -85,7 +85,7 @@ class Flatten(DnnLayer, ExplicitInOutSize):
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(bu.math.flatten, start_axis=start_axis, end_axis=end_axis),
+      y = jax.eval_shape(functools.partial(u.math.flatten, start_axis=start_axis, end_axis=end_axis),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
@@ -101,7 +101,7 @@ class Flatten(DnnLayer, ExplicitInOutSize):
         start_axis = self.start_axis + dim_diff
       else:
         start_axis = x.ndim + self.start_axis
-    return bu.math.flatten(x, start_axis, self.end_axis)
+    return u.math.flatten(x, start_axis, self.end_axis)
 
   def __repr__(self) -> str:
     return f'{self.__class__.__name__}(start_axis={self.start_axis}, end_axis={self.end_axis})'
@@ -153,12 +153,12 @@ class Unflatten(DnnLayer, ExplicitInOutSize):
 
     if in_size is not None:
       self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(bu.math.unflatten, axis=axis, sizes=sizes),
+      y = jax.eval_shape(functools.partial(u.math.unflatten, axis=axis, sizes=sizes),
                          jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
       self.out_size = y.shape
 
   def update(self, x):
-    return bu.math.unflatten(x, self.axis, self.sizes)
+    return u.math.unflatten(x, self.axis, self.sizes)
 
   def __repr__(self):
     return f'{self.__class__.__name__}(axis={self.axis}, sizes={self.sizes})'

brainstate/nn/_rate_rnns.py

@@ -27,7 +27,7 @@ from .. import random, init, functional
 from .._module import Module
 from .._state import ShortTermState, ParamState
 from ..mixin import DelayedInit, Mode
-from ..typing import ArrayLike
+from brainstate.typing import ArrayLike
 
 __all__ = [
   'RNNCell', 'ValinaRNNCell', 'GRUCell', 'MGUCell', 'LSTMCell', 'URLSTMCell',

brainstate/nn/_readout.py

@@ -29,7 +29,7 @@ from ._misc import exp_euler_step
 from .. import environ, init, surrogate
 from .._state import ShortTermState, ParamState
 from ..mixin import Mode
-from ..typing import Size, ArrayLike, DTypeLike
+from brainstate.typing import Size, ArrayLike, DTypeLike
 
 __all__ = [
   'LeakyRateReadout',