brainstate 0.1.4__py2.py3-none-any.whl → 0.1.6__py2.py3-none-any.whl

brainstate/nn/_elementwise.py

@@ -61,7 +61,7 @@ class Threshold(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate
         >>> m = nn.Threshold(0.1, 20)
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -97,7 +97,7 @@ class ReLU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.ReLU()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -106,7 +106,7 @@ class ReLU(ElementWiseBlock):
       An implementation of CReLU - https://arxiv.org/abs/1603.05201
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.ReLU()
         >>> x = random.randn(2).unsqueeze(0)
         >>> output = jax.numpy.concat((m(x), m(-x)))
@@ -151,7 +151,7 @@ class RReLU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.RReLU(0.1, 0.3)
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -205,7 +205,7 @@ class Hardtanh(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Hardtanh(-2, 2)
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -244,7 +244,7 @@ class ReLU6(Hardtanh, ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.ReLU6()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -269,7 +269,7 @@ class Sigmoid(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Sigmoid()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -299,7 +299,7 @@ class Hardsigmoid(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Hardsigmoid()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -325,7 +325,7 @@ class Tanh(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Tanh()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -386,7 +386,7 @@ class Mish(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Mish()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -418,7 +418,7 @@ class Hardswish(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Hardswish()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -452,7 +452,7 @@ class ELU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.ELU()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -489,7 +489,7 @@ class CELU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.CELU()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -530,7 +530,7 @@ class SELU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.SELU()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -559,7 +559,7 @@ class GLU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.GLU()
         >>> x = random.randn(4, 2)
         >>> output = m(x)
@@ -600,7 +600,7 @@ class GELU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.GELU()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -642,7 +642,7 @@ class Hardshrink(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Hardshrink()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -689,7 +689,7 @@ class LeakyReLU(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.LeakyReLU(0.1)
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -721,7 +721,7 @@ class LogSigmoid(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.LogSigmoid()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -749,7 +749,7 @@ class Softplus(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Softplus()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -781,7 +781,7 @@ class Softshrink(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Softshrink()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -843,9 +843,9 @@ class PReLU(ElementWiseBlock):
 
     Examples::
 
-        >>> import brainstate as bst
-        >>> m = bst.nn.PReLU()
-        >>> x = bst.random.randn(2)
+        >>> import brainstate as brainstate
+        >>> m = brainstate.nn.PReLU()
+        >>> x = brainstate.random.randn(2)
         >>> output = m(x)
     """
     __module__ = 'brainstate.nn'
@@ -876,7 +876,7 @@ class Softsign(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Softsign()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -900,7 +900,7 @@ class Tanhshrink(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Tanhshrink()
         >>> x = random.randn(2)
         >>> output = m(x)
@@ -937,7 +937,7 @@ class Softmin(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Softmin(dim=1)
         >>> x = random.randn(2, 3)
         >>> output = m(x)
@@ -990,7 +990,7 @@ class Softmax(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Softmax(dim=1)
         >>> x = random.randn(2, 3)
         >>> output = m(x)
@@ -1027,7 +1027,7 @@ class Softmax2d(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.Softmax2d()
         >>> # you softmax over the 2nd dimension
         >>> x = random.randn(2, 3, 12, 13)
@@ -1062,7 +1062,7 @@ class LogSoftmax(ElementWiseBlock):
     Examples::
 
         >>> import brainstate.nn as nn
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> m = nn.LogSoftmax(dim=1)
         >>> x = random.randn(2, 3)
         >>> output = m(x)

brainstate/nn/_exp_euler.py

@@ -69,27 +69,24 @@ def exp_euler_step(
             f'The input data type should be float64, float32, float16, or bfloat16 '
             f'when using Exponential Euler method. But we got {args[0].dtype}.'
         )
+
+    # drift
     dt = environ.get('dt')
     linear, derivative = vector_grad(fn, argnums=0, return_value=True)(*args, **kwargs)
     linear = u.Quantity(u.get_mantissa(linear), u.get_unit(derivative) / u.get_unit(linear))
     phi = u.math.exprel(dt * linear)
     x_next = args[0] + dt * phi * derivative
 
+    # diffusion
     if diffusion is not None:
-        # unit checking
-        diffusion = diffusion(*args, **kwargs)
-        time_unit = u.get_unit(dt)
-        drift_unit = u.get_unit(derivative)
-        diffusion_unit = u.get_unit(diffusion)
-        # if drift_unit.is_unitless:
-        #   assert diffusion_unit.is_unitless, 'The diffusion term should be unitless when the drift term is unitless.'
-        # else:
-        #   u.fail_for_dimension_mismatch(
-        #     drift_unit, diffusion_unit * time_unit ** 0.5,
-        #     "Drift unit is {drift}, diffusion unit is {diffusion}, ",
-        #     drift=drift_unit, diffusion=diffusion_unit * time_unit ** 0.5
-        #   )
-
-        # diffusion
-        x_next += diffusion * u.math.sqrt(dt) * random.randn_like(args[0])
+        diffusion_part = diffusion(*args, **kwargs) * u.math.sqrt(dt) * random.randn_like(args[0])
+        if u.get_dim(x_next) != u.get_dim(diffusion_part):
+            drift_unit = u.get_unit(x_next)
+            time_unit = u.get_unit(dt)
+            raise ValueError(
+                f"Drift unit is {drift_unit}, "
+                f"expected diffusion unit is {drift_unit / time_unit ** 0.5}, "
+                f"but we got {u.get_unit(diffusion_part)}."
+            )
+        x_next += diffusion_part
     return x_next

brainstate/nn/_inputs.py

@@ -547,7 +547,7 @@ def poisson_input(
                 num_input,
                 p,
                 tar[indices].shape,
-                # check_valid=False,
+                check_valid=False,
                 dtype=tar.dtype
             ),
             tar_val,

brainstate/nn/_linear.py

@@ -361,9 +361,9 @@ class LoRA(Module):
 
     Example usage::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax, jax.numpy as jnp
-        >>> layer = bst.nn.LoRA(3, 2, 4)
+        >>> layer = brainstate.nn.LoRA(3, 2, 4)
         >>> layer.weight.value
     {'lora_a': Array([[ 0.25141352, -0.09826107],
             [ 0.2328382 ,  0.38869813],
@@ -371,8 +371,8 @@ class LoRA(Module):
      'lora_b': Array([[-0.8372317 ,  0.21012013, -0.52999765, -0.31939325],
             [ 0.64234126, -0.42980042,  1.2549229 , -0.47134295]],      dtype=float32)}
         >>> # Wrap around existing layer
-        >>> linear = bst.nn.Linear(3, 4)
-        >>> wrapper = bst.nn.LoRA(3, 2, 4, base_module=linear)
+        >>> linear = brainstate.nn.Linear(3, 4)
+        >>> wrapper = brainstate.nn.LoRA(3, 2, 4, base_module=linear)
         >>> assert wrapper.base_module == linear
         >>> y = layer(jnp.ones((16, 3)))
         >>> y.shape

brainstate/nn/_module.py

@@ -128,9 +128,9 @@ class Module(Node, ParamDesc):
         Examples
         --------
 
-        >>> import brainstate as bst
-        >>> x = bst.random.rand((10, 10))
-        >>> l = bst.nn.Dropout(0.5)
+        >>> import brainstate as brainstate
+        >>> x = brainstate.random.rand((10, 10))
+        >>> l = brainstate.nn.Dropout(0.5)
         >>> y = x >> l
         """
         return self.__call__(other)
@@ -230,7 +230,7 @@ class Module(Node, ParamDesc):
         pass
 
     def __pretty_repr_item__(self, name, value):
-        if name in ['_in_size', '_out_size', '_name']:
+        if name.startswith('_'):
             return None if value is None else (name[1:], value)  # skip the first `_`
         return name, value
 
@@ -266,14 +266,14 @@ class Sequential(Module):
     --------
 
     >>> import jax
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import brainstate.nn as nn
     >>>
     >>> # composing ANN models
     >>> l = nn.Sequential(nn.Linear(100, 10),
     >>>                   jax.nn.relu,
     >>>                   nn.Linear(10, 2))
-    >>> l(bst.random.random((256, 100)))
+    >>> l(brainstate.random.random((256, 100)))
 
     Args:
       modules_as_tuple: The children modules.

brainstate/nn/_module_test.py

@@ -77,7 +77,7 @@ class TestDelay(unittest.TestCase):
             self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
             self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
             self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-        # bst.util.clear_buffer_memory()
+        # brainstate.util.clear_buffer_memory()
 
     def test_jit_erro(self):
         rotation_delay = brainstate.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')

brainstate/nn/_normalizations.py

@@ -488,9 +488,9 @@ class LayerNorm(Module):
 
     Example usage::
 
-      >>> import brainstate as bst
-      >>> x = bst.random.normal(size=(3, 4, 5, 6))
-      >>> layer = bst.nn.LayerNorm(x.shape)
+      >>> import brainstate as brainstate
+      >>> x = brainstate.random.normal(size=(3, 4, 5, 6))
+      >>> layer = brainstate.nn.LayerNorm(x.shape)
       >>> layer.states()
       >>> y = layer(x)
 
@@ -616,9 +616,9 @@ class RMSNorm(Module):
 
     Example usage::
 
-      >>> import brainstate as bst
-      >>> x = bst.random.normal(size=(5, 6))
-      >>> layer = bst.nn.RMSNorm(num_features=6)
+      >>> import brainstate as brainstate
+      >>> x = brainstate.random.normal(size=(5, 6))
+      >>> layer = brainstate.nn.RMSNorm(num_features=6)
       >>> layer.states()
       >>> y = layer(x)
 
@@ -739,14 +739,14 @@ class GroupNorm(Module):
     Example usage::
 
       >>> import numpy as np
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       ...
-      >>> x = bst.random.normal(size=(3, 4, 5, 6))
-      >>> layer = bst.nn.GroupNorm(x.shape, num_groups=3)
+      >>> x = brainstate.random.normal(size=(3, 4, 5, 6))
+      >>> layer = brainstate.nn.GroupNorm(x.shape, num_groups=3)
       >>> layer.states()
       >>> y = layer(x)
-      >>> y = bst.nn.GroupNorm(x.shape, num_groups=1)(x)
-      >>> y2 = bst.nn.LayerNorm(x.shape, reduction_axes=(1, 2, 3))(x)
+      >>> y = brainstate.nn.GroupNorm(x.shape, num_groups=1)(x)
+      >>> y2 = brainstate.nn.LayerNorm(x.shape, reduction_axes=(1, 2, 3))(x)
       >>> np.testing.assert_allclose(y, y2)
 
     Attributes:

brainstate/nn/_normalizations_test.py

@@ -51,21 +51,21 @@ class Test_Normalization(parameterized.TestCase):
     #   normalized_shape=(10, [5, 10])
     # )
     # def test_LayerNorm(self, normalized_shape):
-    #   net = bst.nn.LayerNorm(normalized_shape, )
-    #   input = bst.random.randn(20, 5, 10)
+    #   net = brainstate.nn.LayerNorm(normalized_shape, )
+    #   input = brainstate.random.randn(20, 5, 10)
     #   output = net(input)
     #
     # @parameterized.product(
     #   num_groups=[1, 2, 3, 6]
     # )
     # def test_GroupNorm(self, num_groups):
-    #   input = bst.random.randn(20, 10, 10, 6)
-    #   net = bst.nn.GroupNorm(num_groups=num_groups, num_channels=6, )
+    #   input = brainstate.random.randn(20, 10, 10, 6)
+    #   net = brainstate.nn.GroupNorm(num_groups=num_groups, num_channels=6, )
     #   output = net(input)
     #
     # def test_InstanceNorm(self):
-    #   input = bst.random.randn(20, 10, 10, 6)
-    #   net = bst.nn.InstanceNorm(num_channels=6, )
+    #   input = brainstate.random.randn(20, 10, 10, 6)
+    #   net = brainstate.nn.InstanceNorm(num_channels=6, )
     #   output = net(input)
 
 

brainstate/nn/_poolings.py

@@ -55,8 +55,8 @@ class Flatten(Module):
         end_axis: last dim to flatten (default = -1).
 
     Examples::
-        >>> import brainstate as bst
-        >>> inp = bst.random.randn(32, 1, 5, 5)
+        >>> import brainstate as brainstate
+        >>> inp = brainstate.random.randn(32, 1, 5, 5)
         >>> # With default parameters
         >>> m = Flatten()
         >>> output = m(inp)
@@ -334,10 +334,10 @@ class MaxPool1d(_MaxPool):
 
     Examples::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # pool of size=3, stride=2
         >>> m = MaxPool1d(3, stride=2, channel_axis=-1)
-        >>> input = bst.random.randn(20, 50, 16)
+        >>> input = brainstate.random.randn(20, 50, 16)
         >>> output = m(input)
         >>> output.shape
         (20, 24, 16)
@@ -418,12 +418,12 @@ class MaxPool2d(_MaxPool):
 
     Examples::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # pool of square window of size=3, stride=2
         >>> m = MaxPool2d(3, stride=2)
         >>> # pool of non-square window
         >>> m = MaxPool2d((3, 2), stride=(2, 1), channel_axis=-1)
-        >>> input = bst.random.randn(20, 50, 32, 16)
+        >>> input = brainstate.random.randn(20, 50, 32, 16)
         >>> output = m(input)
         >>> output.shape
         (20, 24, 31, 16)
@@ -509,12 +509,12 @@ class MaxPool3d(_MaxPool):
 
     Examples::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # pool of square window of size=3, stride=2
         >>> m = MaxPool3d(3, stride=2)
         >>> # pool of non-square window
         >>> m = MaxPool3d((3, 2, 2), stride=(2, 1, 2), channel_axis=-1)
-        >>> input = bst.random.randn(20, 50, 44, 31, 16)
+        >>> input = brainstate.random.randn(20, 50, 44, 31, 16)
         >>> output = m(input)
         >>> output.shape
         (20, 24, 43, 15, 16)
@@ -588,10 +588,10 @@ class AvgPool1d(_AvgPool):
 
     Examples::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # pool with window of size=3, stride=2
         >>> m = AvgPool1d(3, stride=2)
-        >>> input = bst.random.randn(20, 50, 16)
+        >>> input = brainstate.random.randn(20, 50, 16)
         >>> m(input).shape
         (20, 24, 16)
 
@@ -665,12 +665,12 @@ class AvgPool2d(_AvgPool):
 
     Examples::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # pool of square window of size=3, stride=2
         >>> m = AvgPool2d(3, stride=2)
         >>> # pool of non-square window
         >>> m = AvgPool2d((3, 2), stride=(2, 1))
-        >>> input = bst.random.randn(20, 50, 32, , 16)
+        >>> input = brainstate.random.randn(20, 50, 32, , 16)
         >>> output = m(input)
         >>> output.shape
         (20, 24, 31, 16)
@@ -753,12 +753,12 @@ class AvgPool3d(_AvgPool):
 
     Examples::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # pool of square window of size=3, stride=2
         >>> m = AvgPool3d(3, stride=2)
         >>> # pool of non-square window
         >>> m = AvgPool3d((3, 2, 2), stride=(2, 1, 2))
-        >>> input = bst.random.randn(20, 50, 44, 31, 16)
+        >>> input = brainstate.random.randn(20, 50, 44, 31, 16)
         >>> output = m(input)
         >>> output.shape
         (20, 24, 43, 15, 16)
@@ -931,10 +931,10 @@ class AdaptiveAvgPool1d(_AdaptivePool):
 
     Examples:
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # target output size of 5
         >>> m = AdaptiveMaxPool1d(5)
-        >>> input = bst.random.randn(1, 64, 8)
+        >>> input = brainstate.random.randn(1, 64, 8)
         >>> output = m(input)
         >>> output.shape
         (1, 5, 8)
@@ -979,22 +979,22 @@ class AdaptiveAvgPool2d(_AdaptivePool):
 
     Examples:
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # target output size of 5x7
         >>> m = AdaptiveMaxPool2d((5, 7))
-        >>> input = bst.random.randn(1, 8, 9, 64)
+        >>> input = brainstate.random.randn(1, 8, 9, 64)
         >>> output = m(input)
         >>> output.shape
         (1, 5, 7, 64)
         >>> # target output size of 7x7 (square)
         >>> m = AdaptiveMaxPool2d(7)
-        >>> input = bst.random.randn(1, 10, 9, 64)
+        >>> input = brainstate.random.randn(1, 10, 9, 64)
         >>> output = m(input)
         >>> output.shape
         (1, 7, 7, 64)
         >>> # target output size of 10x7
         >>> m = AdaptiveMaxPool2d((None, 7))
-        >>> input = bst.random.randn(1, 10, 9, 64)
+        >>> input = brainstate.random.randn(1, 10, 9, 64)
         >>> output = m(input)
         >>> output.shape
         (1, 10, 7, 64)
@@ -1040,22 +1040,22 @@ class AdaptiveAvgPool3d(_AdaptivePool):
 
     Examples:
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> # target output size of 5x7x9
         >>> m = AdaptiveMaxPool3d((5, 7, 9))
-        >>> input = bst.random.randn(1, 8, 9, 10, 64)
+        >>> input = brainstate.random.randn(1, 8, 9, 10, 64)
         >>> output = m(input)
         >>> output.shape
         (1, 5, 7, 9, 64)
         >>> # target output size of 7x7x7 (cube)
         >>> m = AdaptiveMaxPool3d(7)
-        >>> input = bst.random.randn(1, 10, 9, 8, 64)
+        >>> input = brainstate.random.randn(1, 10, 9, 8, 64)
         >>> output = m(input)
         >>> output.shape
         (1, 7, 7, 7, 64)
         >>> # target output size of 7x9x8
         >>> m = AdaptiveMaxPool3d((7, None, None))
-        >>> input = bst.random.randn(1, 10, 9, 8, 64)
+        >>> input = brainstate.random.randn(1, 10, 9, 8, 64)
         >>> output = m(input)
         >>> output.shape
         (1, 7, 9, 8, 64)

brainstate/nn/_synapse.py

@@ -123,9 +123,6 @@ class Expon(Synapse, AlignPost):
         g = exp_euler_step(lambda g: self.sum_current_inputs(-g) / self.tau, self.g.value)
         self.g.value = self.sum_delta_inputs(g)
         if x is not None: self.g.value += x
-        return self.update_return()
-
-    def update_return(self) -> PyTree:
         return self.g.value
 
 
@@ -232,9 +229,6 @@ class DualExpon(Synapse, AlignPost):
         if x is not None:
             self.g_rise.value += x
             self.g_decay.value += x
-        return self.update_return()
-
-    def update_return(self) -> PyTree:
         return self.a * (self.g_decay.value - self.g_rise.value)
 
 
@@ -414,12 +408,8 @@ class AMPA(Synapse):
         t = environ.get('t')
         self.spike_arrival_time.value = u.math.where(pre_spike, t, self.spike_arrival_time.value)
         TT = ((t - self.spike_arrival_time.value) < self.T_duration) * self.T
-        dg = lambda g: self.alpha * TT * (1 - g) - self.beta * g
+        dg = lambda g: self.alpha * TT * (1 * u.get_unit(g) - g) - self.beta * g
         self.g.value = exp_euler_step(dg, self.g.value)
-        return self.update_return()
-
-    def update_return(self) -> PyTree:
-        """Return the synaptic conductance value."""
         return self.g.value
 
 
@@ -513,4 +503,3 @@ class GABAa(AMPA):
             in_size=in_size,
             g_initializer=g_initializer
         )
-

brainstate/nn/_utils.py

@@ -62,7 +62,7 @@ def count_parameters(
 
     Parameters:
     -----------
-    model : bst.nn.Module
+    model : brainstate.nn.Module
         The neural network model for which to count parameters.
 
     Returns: