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

brainstate/nn/metrics.py

@@ -61,12 +61,12 @@ class Average(Metric):
     Example usage::
 
       >>> import jax.numpy as jnp
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
 
       >>> batch_loss = jnp.array([1, 2, 3, 4])
       >>> batch_loss2 = jnp.array([3, 2, 1, 0])
 
-      >>> metrics = bst.nn.metrics.Average()
+      >>> metrics = brainstate.nn.metrics.Average()
       >>> metrics.compute()
       Array(nan, dtype=float32)
       >>> metrics.update(values=batch_loss)
@@ -223,7 +223,7 @@ class Accuracy(Average):
 
     Example usage::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax, jax.numpy as jnp
 
       >>> logits = jax.random.normal(jax.random.key(0), (5, 2))
@@ -231,7 +231,7 @@ class Accuracy(Average):
       >>> logits2 = jax.random.normal(jax.random.key(1), (5, 2))
       >>> labels2 = jnp.array([0, 1, 1, 1, 1])
 
-      >>> metrics = bst.nn.metrics.Accuracy()
+      >>> metrics = brainstate.nn.metrics.Accuracy()
       >>> metrics.compute()
       Array(nan, dtype=float32)
       >>> metrics.update(logits=logits, labels=labels)

brainstate/optim/_optax_optimizer.py

@@ -36,29 +36,29 @@ class OptaxOptimizer(Optimizer):
 
       >>> import jax
       >>> import jax.numpy as jnp
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import optax
       ...
-      >>> class Model(bst.nn.Module):
+      >>> class Model(brainstate.nn.Module):
       ...   def __init__(self):
       ...     super().__init__()
-      ...     self.linear1 = bst.nn.Linear(2, 3)
-      ...     self.linear2 = bst.nn.Linear(3, 4)
+      ...     self.linear1 = brainstate.nn.Linear(2, 3)
+      ...     self.linear2 = brainstate.nn.Linear(3, 4)
       ...   def __call__(self, x):
       ...     return self.linear2(self.linear1(x))
       ...
-      >>> x = bst.random.randn(1, 2)
+      >>> x = brainstate.random.randn(1, 2)
       >>> y = jnp.ones((1, 4))
       ...
       >>> model = Model()
       >>> tx = optax.adam(1e-3)
-      >>> optimizer = bst.optim.OptaxOptimizer(tx)
-      >>> optimizer.register_trainable_weights(model.states(bst.ParamState))
+      >>> optimizer = brainstate.optim.OptaxOptimizer(tx)
+      >>> optimizer.register_trainable_weights(model.states(brainstate.ParamState))
       ...
       >>> loss_fn = lambda: ((model(x) - y) ** 2).mean()
       >>> loss_fn()
       Array(1.7055722, dtype=float32)
-      >>> grads = bst.augment.grad(loss_fn, model.states(bst.ParamState))()
+      >>> grads = brainstate.augment.grad(loss_fn, model.states(brainstate.ParamState))()
       >>> optimizer.update(grads)
       >>> loss_fn()
       Array(1.6925814, dtype=float32)

brainstate/random/_rand_funs.py

@@ -78,8 +78,8 @@ def rand(*dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
 
     Examples
     --------
-    >>> import brainstate as bst
-    >>> bst.random.rand(3,2)
+    >>> import brainstate as brainstate
+    >>> brainstate.random.rand(3,2)
     array([[ 0.14022471,  0.96360618],  #random
            [ 0.37601032,  0.25528411],  #random
            [ 0.49313049,  0.94909878]]) #random
@@ -135,31 +135,31 @@ def randint(low, high=None, size: Optional[Size] = None,
 
     Examples
     --------
-    >>> import brainstate as bst
-    >>> bst.random.randint(2, size=10)
+    >>> import brainstate as brainstate
+    >>> brainstate.random.randint(2, size=10)
     array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0]) # random
-    >>> bst.random.randint(1, size=10)
+    >>> brainstate.random.randint(1, size=10)
     array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
 
     Generate a 2 x 4 array of ints between 0 and 4, inclusive:
 
-    >>> bst.random.randint(5, size=(2, 4))
+    >>> brainstate.random.randint(5, size=(2, 4))
     array([[4, 0, 2, 1], # random
            [3, 2, 2, 0]])
 
     Generate a 1 x 3 array with 3 different upper bounds
 
-    >>> bst.random.randint(1, [3, 5, 10])
+    >>> brainstate.random.randint(1, [3, 5, 10])
     array([2, 2, 9]) # random
 
     Generate a 1 by 3 array with 3 different lower bounds
 
-    >>> bst.random.randint([1, 5, 7], 10)
+    >>> brainstate.random.randint([1, 5, 7], 10)
     array([9, 8, 7]) # random
 
     Generate a 2 by 4 array using broadcasting with dtype of uint8
 
-    >>> bst.random.randint([1, 3, 5, 7], [[10], [20]], dtype=np.uint8)
+    >>> brainstate.random.randint([1, 3, 5, 7], [[10], [20]], dtype=np.uint8)
     array([[ 8,  6,  9,  7], # random
            [ 1, 16,  9, 12]], dtype=uint8)
     """
@@ -219,12 +219,12 @@ def random_integers(low,
 
     Examples
     --------
-    >>> import brainstate as bst
-    >>> bst.random.random_integers(5)
+    >>> import brainstate as brainstate
+    >>> brainstate.random.random_integers(5)
     4 # random
-    >>> type(bst.random.random_integers(5))
+    >>> type(brainstate.random.random_integers(5))
     <class 'numpy.int64'>
-    >>> bst.random.random_integers(5, size=(3,2))
+    >>> brainstate.random.random_integers(5, size=(3,2))
     array([[5, 4], # random
            [3, 3],
            [4, 5]])
@@ -233,13 +233,13 @@ def random_integers(low,
     numbers between 0 and 2.5, inclusive (*i.e.*, from the set
     :math:`{0, 5/8, 10/8, 15/8, 20/8}`):
 
-    >>> 2.5 * (bst.random.random_integers(5, size=(5,)) - 1) / 4.
+    >>> 2.5 * (brainstate.random.random_integers(5, size=(5,)) - 1) / 4.
     array([ 0.625,  1.25 ,  0.625,  0.625,  2.5  ]) # random
 
     Roll two six sided dice 1000 times and sum the results:
 
-    >>> d1 = bst.random.random_integers(1, 6, 1000)
-    >>> d2 = bst.random.random_integers(1, 6, 1000)
+    >>> d1 = brainstate.random.random_integers(1, 6, 1000)
+    >>> d2 = brainstate.random.random_integers(1, 6, 1000)
     >>> dsums = d1 + d2
 
     Display results as a histogram:
@@ -301,13 +301,13 @@ def randn(*dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
 
     Examples
     --------
-    >>> import brainstate as bst
-    >>> bst.random.randn()
+    >>> import brainstate as brainstate
+    >>> brainstate.random.randn()
     2.1923875335537315  # random
 
     Two-by-four array of samples from N(3, 6.25):
 
-    >>> 3 + 2.5 * bst.random.randn(2, 4)
+    >>> 3 + 2.5 * brainstate.random.randn(2, 4)
     array([[-4.49401501,  4.00950034, -1.81814867,  7.29718677],   # random
            [ 0.39924804,  4.68456316,  4.99394529,  4.84057254]])  # random
     """
@@ -359,17 +359,17 @@ def random_sample(size: Optional[Size] = None, key: Optional[SeedOrKey] = None,
 
     Examples
     --------
-    >>> import brainstate as bst
-    >>> bst.random.random_sample()
+    >>> import brainstate as brainstate
+    >>> brainstate.random.random_sample()
     0.47108547995356098 # random
-    >>> type(bst.random.random_sample())
+    >>> type(brainstate.random.random_sample())
     <class 'float'>
-    >>> bst.random.random_sample((5,))
+    >>> brainstate.random.random_sample((5,))
     array([ 0.30220482,  0.86820401,  0.1654503 ,  0.11659149,  0.54323428]) # random
 
     Three-by-two array of random numbers from [-5, 0):
 
-    >>> 5 * bst.random.random_sample((3, 2)) - 5
+    >>> 5 * brainstate.random.random_sample((3, 2)) - 5
     array([[-3.99149989, -0.52338984], # random
            [-2.99091858, -0.79479508],
            [-1.23204345, -1.75224494]])
@@ -450,34 +450,34 @@ def choice(a, size: Optional[Size] = None, replace=True, p=None,
     --------
     Generate a uniform random sample from np.arange(5) of size 3:
 
-    >>> import brainstate as bst
-    >>> bst.random.choice(5, 3)
+    >>> import brainstate as brainstate
+    >>> brainstate.random.choice(5, 3)
     array([0, 3, 4]) # random
     >>> #This is equivalent to brainpy.math.random.randint(0,5,3)
 
     Generate a non-uniform random sample from np.arange(5) of size 3:
 
-    >>> bst.random.choice(5, 3, p=[0.1, 0, 0.3, 0.6, 0])
+    >>> brainstate.random.choice(5, 3, p=[0.1, 0, 0.3, 0.6, 0])
     array([3, 3, 0]) # random
 
     Generate a uniform random sample from np.arange(5) of size 3 without
     replacement:
 
-    >>> bst.random.choice(5, 3, replace=False)
+    >>> brainstate.random.choice(5, 3, replace=False)
     array([3,1,0]) # random
     >>> #This is equivalent to brainpy.math.random.permutation(np.arange(5))[:3]
 
     Generate a non-uniform random sample from np.arange(5) of size
     3 without replacement:
 
-    >>> bst.random.choice(5, 3, replace=False, p=[0.1, 0, 0.3, 0.6, 0])
+    >>> brainstate.random.choice(5, 3, replace=False, p=[0.1, 0, 0.3, 0.6, 0])
     array([2, 3, 0]) # random
 
     Any of the above can be repeated with an arbitrary array-like
     instead of just integers. For instance:
 
     >>> aa_milne_arr = ['pooh', 'rabbit', 'piglet', 'Christopher']
-    >>> bst.random.choice(aa_milne_arr, 5, p=[0.5, 0.1, 0.1, 0.3])
+    >>> brainstate.random.choice(aa_milne_arr, 5, p=[0.5, 0.1, 0.1, 0.3])
     array(['pooh', 'pooh', 'pooh', 'Christopher', 'piglet'], # random
           dtype='<U11')
     """
@@ -519,15 +519,15 @@ def permutation(x,
 
     Examples
     --------
-    >>> import brainstate as bst
-    >>> bst.random.permutation(10)
+    >>> import brainstate as brainstate
+    >>> brainstate.random.permutation(10)
     array([1, 7, 4, 3, 0, 9, 2, 5, 8, 6]) # random
 
-    >>> bst.random.permutation([1, 4, 9, 12, 15])
+    >>> brainstate.random.permutation([1, 4, 9, 12, 15])
     array([15,  1,  9,  4, 12]) # random
 
     >>> arr = np.arange(9).reshape((3, 3))
-    >>> bst.random.permutation(arr)
+    >>> brainstate.random.permutation(arr)
     array([[6, 7, 8], # random
            [0, 1, 2],
            [3, 4, 5]])
@@ -557,16 +557,16 @@ def shuffle(x, axis=0, key: Optional[SeedOrKey] = None):
 
     Examples
     --------
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> arr = np.arange(10)
-    >>> bst.random.shuffle(arr)
+    >>> brainstate.random.shuffle(arr)
     >>> arr
     [1 7 5 2 9 4 3 6 0 8] # random
 
     Multi-dimensional arrays are only shuffled along the first axis:
 
     >>> arr = np.arange(9).reshape((3, 3))
-    >>> bst.random.shuffle(arr)
+    >>> brainstate.random.shuffle(arr)
     >>> arr
     array([[3, 4, 5], # random
            [6, 7, 8],

brainstate/random/_rand_funs_test.py

@@ -562,6 +562,6 @@ class TestRandom(unittest.TestCase):
 
 # class TestRandomKey(unittest.TestCase):
 #   def test_clear_memory(self):
-#     bst.random.split_key()
-#     print(bst.random.DEFAULT.value)
-#     self.assertTrue(isinstance(bst.random.DEFAULT.value, np.ndarray))
+#     brainstate.random.split_key()
+#     print(brainstate.random.DEFAULT.value)
+#     self.assertTrue(isinstance(brainstate.random.DEFAULT.value, np.ndarray))

brainstate/random/_rand_seed.py

@@ -183,16 +183,16 @@ def seed_context(seed_or_key: SeedOrKey):
 
     Examples:
 
-    >>> import brainstate as bst
-    >>> print(bst.random.rand(2))
+    >>> import brainstate as brainstate
+    >>> print(brainstate.random.rand(2))
     [0.57721865 0.9820676 ]
-    >>> print(bst.random.rand(2))
+    >>> print(brainstate.random.rand(2))
     [0.8511752  0.95312667]
-    >>> with bst.random.seed_context(42):
-    ...     print(bst.random.rand(2))
+    >>> with brainstate.random.seed_context(42):
+    ...     print(brainstate.random.rand(2))
     [0.95598125 0.4032725 ]
-    >>> with bst.random.seed_context(42):
-    ...     print(bst.random.rand(2))
+    >>> with brainstate.random.seed_context(42):
+    ...     print(brainstate.random.rand(2))
     [0.95598125 0.4032725 ]
 
     Args:

brainstate/random/_rand_state.py

@@ -384,7 +384,10 @@ class RandomState(State):
         loc = _check_py_seq(loc)
         scale = _check_py_seq(scale)
         if size is None:
-            size = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
+            size = lax.broadcast_shapes(
+                jnp.shape(loc) if loc is not None else (),
+                jnp.shape(scale) if scale is not None else ()
+            )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = _loc_scale(loc, scale, jr.logistic(key, shape=_size2shape(size), dtype=dtype))
@@ -399,7 +402,10 @@ class RandomState(State):
         loc = _check_py_seq(loc)
         scale = _check_py_seq(scale)
         if size is None:
-            size = lax.broadcast_shapes(jnp.shape(scale), jnp.shape(loc))
+            size = lax.broadcast_shapes(
+                jnp.shape(scale) if scale is not None else (),
+                jnp.shape(loc) if loc is not None else ()
+            )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = _loc_scale(loc, scale, jr.normal(key, shape=_size2shape(size), dtype=dtype))
@@ -456,7 +462,7 @@ class RandomState(State):
                        dtype: DTypeLike = None):
         shape = _check_py_seq(shape)
         if size is None:
-            size = jnp.shape(shape)
+            size = jnp.shape(shape) if shape is not None else ()
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = jr.gamma(key, a=shape, shape=_size2shape(size), dtype=dtype)
@@ -477,7 +483,7 @@ class RandomState(State):
                    dtype: DTypeLike = None):
         df = _check_py_seq(df)
         if size is None:
-            size = jnp.shape(size)
+            size = jnp.shape(size) if size is not None else ()
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
         r = jr.t(key, df=df, shape=_size2shape(size), dtype=dtype)
@@ -606,8 +612,8 @@ class RandomState(State):
 
         if size is None:
             size = jnp.broadcast_shapes(
-                jnp.shape(mean),
-                jnp.shape(sigma)
+                jnp.shape(mean) if mean is not None else (),
+                jnp.shape(sigma) if sigma is not None else ()
             )
         key = self.split_key() if key is None else _formalize_key(key)
         dtype = dtype or environ.dftype()
@@ -822,7 +828,7 @@ class RandomState(State):
         a = _check_py_seq(a)
         scale = _check_py_seq(scale)
         if size is None:
-            size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale))
+            size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale) if scale is not None else ())
         else:
             if jnp.size(a) > 1:
                 raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')

brainstate/surrogate.py

@@ -106,11 +106,11 @@ class Surrogate(PrettyObject):
     Examples
     --------
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import brainstate.nn as nn
     >>> import jax.numpy as jnp
 
-    >>> class MySurrogate(bst.surrogate.Surrogate):
+    >>> class MySurrogate(brainstate.surrogate.Surrogate):
     ...   def __init__(self, alpha=1.):
     ...     super().__init__()
     ...     self.alpha = alpha
@@ -236,11 +236,11 @@ def sigmoid(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-2, 2, 1000)
        >>> for alpha in [1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.sigmoid)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.sigmoid)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -355,11 +355,11 @@ def piecewise_quadratic(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.piecewise_quadratic)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.piecewise_quadratic)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -522,11 +522,11 @@ def piecewise_exp(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.piecewise_exp)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.piecewise_exp)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -621,11 +621,11 @@ def soft_sign(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.soft_sign)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.soft_sign)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -706,11 +706,11 @@ def arctan(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.arctan)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.arctan)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -804,11 +804,11 @@ def nonzero_sign_log(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.nonzero_sign_log)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.nonzero_sign_log)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -893,11 +893,11 @@ def erf(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.nonzero_sign_log)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.nonzero_sign_log)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -1000,12 +1000,12 @@ def piecewise_leaky_relu(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for c in [0.01, 0.05, 0.1]:
        >>>   for w in [1., 2.]:
-       >>>     grads1 = bst.augment.vector_grad(bst.surrogate.piecewise_leaky_relu)(xs, c=c, w=w)
+       >>>     grads1 = brainstate.augment.vector_grad(brainstate.surrogate.piecewise_leaky_relu)(xs, c=c, w=w)
        >>>     plt.plot(xs, grads1, label=f'x={c}, w={w}')
        >>> plt.legend()
        >>> plt.show()
@@ -1113,12 +1113,12 @@ def squarewave_fourier_series(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for n in [2, 4, 8]:
-       >>>   f = bst.surrogate.SquarewaveFourierSeries(n=n)
-       >>>   grads1 = bst.augment.vector_grad(f)(xs)
+       >>>   f = brainstate.surrogate.SquarewaveFourierSeries(n=n)
+       >>>   grads1 = brainstate.augment.vector_grad(f)(xs)
        >>>   plt.plot(xs, grads1, label=f'n={n}')
        >>> plt.legend()
        >>> plt.show()
@@ -1214,12 +1214,12 @@ def s2nn(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.s2nn)(xs, 4., 1.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.s2nn)(xs, 4., 1.)
        >>> plt.plot(xs, grads, label=r'$\alpha=4, \beta=1$')
-       >>> grads = bst.augment.vector_grad(bst.surrogate.s2nn)(xs, 8., 2.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.s2nn)(xs, 8., 2.)
        >>> plt.plot(xs, grads, label=r'$\alpha=8, \beta=2$')
        >>> plt.legend()
        >>> plt.show()
@@ -1315,11 +1315,11 @@ def q_pseudo_spike(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.q_pseudo_spike)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.q_pseudo_spike)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha=$' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -1413,10 +1413,10 @@ def leaky_relu(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.leaky_relu)(xs, 0., 1.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.leaky_relu)(xs, 0., 1.)
        >>> plt.plot(xs, grads, label=r'$\alpha=0., \beta=1.$')
        >>> plt.legend()
        >>> plt.show()
@@ -1517,10 +1517,10 @@ def log_tailed_relu(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.leaky_relu)(xs, 0., 1.)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.leaky_relu)(xs, 0., 1.)
        >>> plt.plot(xs, grads, label=r'$\alpha=0., \beta=1.$')
        >>> plt.legend()
        >>> plt.show()
@@ -1596,12 +1596,12 @@ def relu_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for s in [0.5, 1.]:
        >>>   for w in [1, 2.]:
-       >>>     grads = bst.augment.vector_grad(bst.surrogate.relu_grad)(xs, s, w)
+       >>>     grads = brainstate.augment.vector_grad(brainstate.surrogate.relu_grad)(xs, s, w)
        >>>     plt.plot(xs, grads, label=r'$\alpha=$' + f'{s}, width={w}')
        >>> plt.legend()
        >>> plt.show()
@@ -1678,11 +1678,11 @@ def gaussian_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for s in [0.5, 1., 2.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.gaussian_grad)(xs, s, 0.5)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.gaussian_grad)(xs, s, 0.5)
        >>>   plt.plot(xs, grads, label=r'$\alpha=0.5, \sigma=$' + str(s))
        >>> plt.legend()
        >>> plt.show()
@@ -1773,10 +1773,10 @@ def multi_gaussian_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
-       >>> grads = bst.augment.vector_grad(bst.surrogate.multi_gaussian_grad)(xs)
+       >>> grads = brainstate.augment.vector_grad(brainstate.surrogate.multi_gaussian_grad)(xs)
        >>> plt.plot(xs, grads)
        >>> plt.show()
 
@@ -1855,11 +1855,11 @@ def inv_square_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-1, 1, 1000)
        >>> for alpha in [1., 10., 100.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.inv_square_grad)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.inv_square_grad)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()
@@ -1929,11 +1929,11 @@ def slayer_grad(
 
        >>> import jax
        >>> import brainstate.nn as nn
-       >>> import brainstate as bst
+       >>> import brainstate as brainstate
        >>> import matplotlib.pyplot as plt
        >>> xs = jax.numpy.linspace(-3, 3, 1000)
        >>> for alpha in [0.5, 1., 2., 4.]:
-       >>>   grads = bst.augment.vector_grad(bst.surrogate.slayer_grad)(xs, alpha)
+       >>>   grads = brainstate.augment.vector_grad(brainstate.surrogate.slayer_grad)(xs, alpha)
        >>>   plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha))
        >>> plt.legend()
        >>> plt.show()