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

brainstate/nn/_dynamics.py

@@ -42,7 +42,7 @@ import numpy as np
 from brainstate import environ
 from brainstate._state import State
 from brainstate.graph import Node
-from brainstate.mixin import ParamDescriber, UpdateReturn
+from brainstate.mixin import ParamDescriber
 from brainstate.typing import Size, ArrayLike, PyTree
 from ._delay import StateWithDelay, Delay
 from ._module import Module
@@ -101,7 +101,7 @@ class Projection(Module):
             raise ValueError('Do not implement the update() function.')
 
 
-class Dynamics(Module, UpdateReturn):
+class Dynamics(Module):
     """
     Base class for implementing neural dynamics models in BrainState.
 
@@ -214,13 +214,13 @@ class Dynamics(Module, UpdateReturn):
         # in-/out- size of neuron population
         self.out_size = self.in_size
 
-    def __pretty_repr_item__(self, name, value):
-        if name in [
-            '_before_updates', '_after_updates', '_current_inputs', '_delta_inputs',
-            '_in_size', '_out_size', '_name', '_mode',
-        ]:
-            return (name, value) if value is None else (name[1:], value)  # skip the first `_`
-        return super().__pretty_repr_item__(name, value)
+    # def __pretty_repr_item__(self, name, value):
+    #     if name in [
+    #         '_before_updates', '_after_updates', '_current_inputs', '_delta_inputs',
+    #         '_in_size', '_out_size', '_name', '_mode',
+    #     ]:
+    #         return (name, value) if value is None else (name[1:], value)  # skip the first `_`
+    #     return super().__pretty_repr_item__(name, value)
 
     @property
     def varshape(self):
@@ -470,21 +470,30 @@ class Dynamics(Module, UpdateReturn):
         if self._current_inputs is None:
             return init
         if label is None:
-            # no label
-            for key in tuple(self._current_inputs.keys()):
-                out = self._current_inputs[key]
-                init = init + (out(*args, **kwargs) if callable(out) else out)
-                if not callable(out):
-                    self._current_inputs.pop(key)
+            filter_fn = lambda k: True
         else:
-            # has label
             label_repr = _input_label_start(label)
-            for key in tuple(self._current_inputs.keys()):
-                if key.startswith(label_repr):
-                    out = self._current_inputs[key]
-                    init = init + (out(*args, **kwargs) if callable(out) else out)
-                    if not callable(out):
-                        self._current_inputs.pop(key)
+            filter_fn = lambda k: k.startswith(label_repr)
+        for key in tuple(self._current_inputs.keys()):
+            if filter_fn(key):
+                out = self._current_inputs[key]
+                if callable(out):
+                    try:
+                        init = init + out(*args, **kwargs)
+                    except Exception as e:
+                        raise ValueError(
+                            f'Error in delta input value {key}: {out}\n'
+                            f'Error: {e}'
+                        ) from e
+                else:
+                    try:
+                        init = init + out
+                    except Exception as e:
+                        raise ValueError(
+                            f'Error in delta input value {key}: {out}\n'
+                            f'Error: {e}'
+                        ) from e
+                    self._current_inputs.pop(key)
         return init
 
     def sum_delta_inputs(
@@ -529,21 +538,30 @@ class Dynamics(Module, UpdateReturn):
         if self._delta_inputs is None:
             return init
         if label is None:
-            # no label
-            for key in tuple(self._delta_inputs.keys()):
-                out = self._delta_inputs[key]
-                init = init + (out(*args, **kwargs) if callable(out) else out)
-                if not callable(out):
-                    self._delta_inputs.pop(key)
+            filter_fn = lambda k: True
         else:
-            # has label
             label_repr = _input_label_start(label)
-            for key in tuple(self._delta_inputs.keys()):
-                if key.startswith(label_repr):
-                    out = self._delta_inputs[key]
-                    init = init + (out(*args, **kwargs) if callable(out) else out)
-                    if not callable(out):
-                        self._delta_inputs.pop(key)
+            filter_fn = lambda k: k.startswith(label_repr)
+        for key in tuple(self._delta_inputs.keys()):
+            if filter_fn(key):
+                out = self._delta_inputs[key]
+                if callable(out):
+                    try:
+                        init = init + out(*args, **kwargs)
+                    except Exception as e:
+                        raise ValueError(
+                            f'Error in delta input function {key}: {out}\n'
+                            f'Error: {e}'
+                        ) from e
+                else:
+                    try:
+                        init = init + out
+                    except Exception as e:
+                        raise ValueError(
+                            f'Error in delta input value {key}: {out}\n'
+                            f'Error: {e}'
+                        ) from e
+                    self._delta_inputs.pop(key)
         return init
 
     @property

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/_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)