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

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:

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: