brainstate 0.2.0__py2.py3-none-any.whl → 0.2.2__py2.py3-none-any.whl

brainstate/nn/_dynamics_test.py

@@ -42,20 +42,6 @@ class TestDynamics(unittest.TestCase):
         with self.assertRaises(ValueError):
             brainstate.nn.Dynamics(in_size="invalid")
 
-    def test_input_handling(self):
-        dyn = brainstate.nn.Dynamics(in_size=10)
-        dyn.add_current_input("test_current", lambda: np.random.rand(10))
-        dyn.add_delta_input("test_delta", lambda: np.random.rand(10))
-
-        self.assertIn("test_current", dyn.current_inputs)
-        self.assertIn("test_delta", dyn.delta_inputs)
-
-    def test_duplicate_input_key(self):
-        dyn = brainstate.nn.Dynamics(in_size=10)
-        dyn.add_current_input("test", lambda: np.random.rand(10))
-        with self.assertRaises(ValueError):
-            dyn.add_current_input("test", lambda: np.random.rand(10))
-
     def test_varshape(self):
         dyn = brainstate.nn.Dynamics(in_size=(2, 3))
         self.assertEqual(dyn.varshape, (2, 3))

brainstate/nn/_embedding.py

@@ -161,8 +161,8 @@ class Embedding(Module):
 
     .. code-block:: python
 
-        >>> import brainstate as bst
-        >>> embedding = bst.nn.Embedding(num_embeddings=10, embedding_size=3)
+        >>> import brainstate as brainstate
+        >>> embedding = brainstate.nn.Embedding(num_embeddings=10, embedding_size=3)
         >>> embedding.weight.value.shape
         (10, 3)
 
@@ -191,7 +191,7 @@ class Embedding(Module):
 
     .. code-block:: python
 
-        >>> embedding = bst.nn.Embedding(num_embeddings=10, embedding_size=3, padding_idx=0)
+        >>> embedding = brainstate.nn.Embedding(num_embeddings=10, embedding_size=3, padding_idx=0)
         >>> # The embedding at index 0 will remain zeros and not be updated during training
         >>> indices = jnp.array([0, 2, 0, 5])
         >>> output = embedding(indices)
@@ -202,7 +202,7 @@ class Embedding(Module):
 
     .. code-block:: python
 
-        >>> embedding = bst.nn.Embedding(num_embeddings=10, embedding_size=3, max_norm=1.0)
+        >>> embedding = brainstate.nn.Embedding(num_embeddings=10, embedding_size=3, max_norm=1.0)
         >>> # All embeddings accessed in a forward pass are renormalized to have norm <= 1.0
 
     Load pretrained embeddings:
@@ -214,7 +214,7 @@ class Embedding(Module):
         >>> pretrained = jnp.array([[1.0, 2.0, 3.0],
         ...                         [4.0, 5.0, 6.0],
         ...                         [7.0, 8.0, 9.0]])
-        >>> embedding = bst.nn.Embedding.from_pretrained(pretrained, param_type=brainstate.FakeState)
+        >>> embedding = brainstate.nn.Embedding.from_pretrained(pretrained, param_type=brainstate.FakeState)
         >>> embedding.weight.value.shape
         (3, 3)
     """
@@ -310,11 +310,11 @@ class Embedding(Module):
         .. code-block:: python
 
             >>> import jax.numpy as jnp
-            >>> import brainstate as bst
+            >>> import brainstate as brainstate
             >>> pretrained = jnp.array([[1.0, 2.0, 3.0],
             ...                         [4.0, 5.0, 6.0],
             ...                         [7.0, 8.0, 9.0]])
-            >>> embedding = bst.nn.Embedding.from_pretrained(pretrained)
+            >>> embedding = brainstate.nn.Embedding.from_pretrained(pretrained)
             >>> embedding.weight.value.shape
             (3, 3)
             >>> indices = jnp.array([1])

brainstate/nn/_exp_euler.py

@@ -112,11 +112,11 @@ def exp_euler_step(
 
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Set time step in environment
-        >>> bst.environ.set(dt=0.01)
+        >>> brainstate.environ.set(dt=0.01)
         >>>
         >>> # Define drift function
         >>> def drift(x, t):
@@ -126,7 +126,7 @@ def exp_euler_step(
         >>> x0 = jnp.array(1.0)
         >>>
         >>> # Single integration step
-        >>> x1 = bst.nn.exp_euler_step(drift, x0, None)
+        >>> x1 = brainstate.nn.exp_euler_step(drift, x0, None)
         >>> print(x1)  # Should be close to exp(-0.01) ≈ 0.99
 
     **SDE Integration:**
@@ -135,11 +135,11 @@ def exp_euler_step(
 
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Set time step
-        >>> bst.environ.set(dt=0.01)
+        >>> brainstate.environ.set(dt=0.01)
         >>>
         >>> # Define drift and diffusion
         >>> theta = 0.5
@@ -155,16 +155,16 @@ def exp_euler_step(
         >>> x0 = jnp.array(1.0)
         >>>
         >>> # Single SDE integration step
-        >>> x1 = bst.nn.exp_euler_step(drift, diffusion, x0, None)
+        >>> x1 = brainstate.nn.exp_euler_step(drift, diffusion, x0, None)
 
     **Multi-dimensional system:**
 
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
-        >>> bst.environ.set(dt=0.01)
+        >>> brainstate.environ.set(dt=0.01)
         >>>
         >>> # Coupled oscillator system
         >>> def drift(x, t):
@@ -172,7 +172,7 @@ def exp_euler_step(
         ...     return jnp.array([-x1 + x2, -x2 - x1])
         >>>
         >>> x0 = jnp.array([1.0, 0.0])
-        >>> x1 = bst.nn.exp_euler_step(drift, x0, None)
+        >>> x1 = brainstate.nn.exp_euler_step(drift, x0, None)
 
     See Also
     --------

brainstate/nn/_linear.py

@@ -76,18 +76,18 @@ class Linear(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a linear layer
-        >>> layer = bst.nn.Linear((10,), (5,))
+        >>> layer = brainstate.nn.Linear((10,), (5,))
         >>> x = jnp.ones((32, 10))
         >>> y = layer(x)
         >>> y.shape
         (32, 5)
         >>>
         >>> # Linear layer without bias
-        >>> layer = bst.nn.Linear((10,), (5,), b_init=None)
+        >>> layer = brainstate.nn.Linear((10,), (5,), b_init=None)
         >>> y = layer(x)
         >>> y.shape
         (32, 5)
@@ -171,11 +171,11 @@ class SignedWLinear(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a signed weight linear layer with all positive weights
-        >>> layer = bst.nn.SignedWLinear((10,), (5,))
+        >>> layer = brainstate.nn.SignedWLinear((10,), (5,))
         >>> x = jnp.ones((32, 10))
         >>> y = layer(x)
         >>> y.shape
@@ -183,7 +183,7 @@ class SignedWLinear(Module):
         >>>
         >>> # With custom sign matrix (e.g., inhibitory connections)
         >>> w_sign = jnp.ones((10, 5)) * -1.0  # all negative
-        >>> layer = bst.nn.SignedWLinear((10,), (5,), w_sign=w_sign)
+        >>> layer = brainstate.nn.SignedWLinear((10,), (5,), w_sign=w_sign)
         >>> y = layer(x)
         >>> y.shape
         (32, 5)
@@ -274,18 +274,18 @@ class ScaledWSLinear(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a weight-standardized linear layer
-        >>> layer = bst.nn.ScaledWSLinear((10,), (5,))
+        >>> layer = brainstate.nn.ScaledWSLinear((10,), (5,))
         >>> x = jnp.ones((32, 10))
         >>> y = layer(x)
         >>> y.shape
         (32, 5)
         >>>
         >>> # Without learnable gain
-        >>> layer = bst.nn.ScaledWSLinear((10,), (5,), ws_gain=False)
+        >>> layer = brainstate.nn.ScaledWSLinear((10,), (5,), ws_gain=False)
         >>> y = layer(x)
         >>> y.shape
         (32, 5)
@@ -375,7 +375,7 @@ class SparseLinear(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import brainunit as u
         >>> import jax.numpy as jnp
         >>>
@@ -384,7 +384,7 @@ class SparseLinear(Module):
         >>> values = jnp.array([1.0, 2.0, 3.0])
         >>> spar_mat = u.sparse.CSR((values, indices[:, 1], indices[:, 0]),
         ...                          shape=(3, 3))
-        >>> layer = bst.nn.SparseLinear(spar_mat, in_size=(3,))
+        >>> layer = brainstate.nn.SparseLinear(spar_mat, in_size=(3,))
         >>> x = jnp.ones((5, 3))
         >>> y = layer(x)
         >>> y.shape
@@ -468,18 +468,18 @@ class AllToAll(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # All-to-all with self-connections
-        >>> layer = bst.nn.AllToAll((10,), (10,), include_self=True)
+        >>> layer = brainstate.nn.AllToAll((10,), (10,), include_self=True)
         >>> x = jnp.ones((32, 10))
         >>> y = layer(x)
         >>> y.shape
         (32, 10)
         >>>
         >>> # All-to-all without self-connections (recurrent layer)
-        >>> layer = bst.nn.AllToAll((10,), (10,), include_self=False)
+        >>> layer = brainstate.nn.AllToAll((10,), (10,), include_self=False)
         >>> y = layer(x)
         >>> y.shape
         (32, 10)
@@ -584,18 +584,18 @@ class OneToOne(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # One-to-one connection
-        >>> layer = bst.nn.OneToOne((10,))
+        >>> layer = brainstate.nn.OneToOne((10,))
         >>> x = jnp.ones((32, 10))
         >>> y = layer(x)
         >>> y.shape
         (32, 10)
         >>>
         >>> # With bias
-        >>> layer = bst.nn.OneToOne((10,), b_init=bst.init.Constant(0.1))
+        >>> layer = brainstate.nn.OneToOne((10,), b_init=brainstate.init.Constant(0.1))
         >>> y = layer(x)
         >>> y.shape
         (32, 10)
@@ -677,19 +677,19 @@ class LoRA(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Standalone LoRA layer
-        >>> layer = bst.nn.LoRA(in_features=10, lora_rank=2, out_features=5)
+        >>> layer = brainstate.nn.LoRA(in_features=10, lora_rank=2, out_features=5)
         >>> x = jnp.ones((32, 10))
         >>> y = layer(x)
         >>> y.shape
         (32, 5)
         >>>
         >>> # Wrap around existing linear layer
-        >>> base = bst.nn.Linear((10,), (5,))
-        >>> lora_layer = bst.nn.LoRA(in_features=10, lora_rank=2,
+        >>> base = brainstate.nn.Linear((10,), (5,))
+        >>> lora_layer = brainstate.nn.LoRA(in_features=10, lora_rank=2,
         ...                           out_features=5, base_module=base)
         >>> y = lora_layer(x)
         >>> y.shape

brainstate/nn/_module.py

@@ -363,22 +363,29 @@ class Sequential(Module):
         self.out_size = in_size
 
     def _format_module(self, module, in_size):
-        if isinstance(module, ParamDescriber):
-            if in_size is None:
-                raise ValueError(
-                    'The input size should be specified. '
-                    f'Please set the in_size attribute of the previous module: \n'
-                    f'{self.layers[-1]}'
-                )
-            module = module(in_size=in_size)
-            assert isinstance(module, Module), 'The module should be an instance of Module.'
-            out_size = module.out_size
-        elif isinstance(module, ElementWiseBlock):
-            out_size = in_size
-        elif isinstance(module, Module):
-            out_size = module.out_size
-        elif callable(module):
-            out_size = in_size
-        else:
-            raise TypeError(f"Unsupported type {type(module)}. ")
+        try:
+            if isinstance(module, ParamDescriber):
+                if in_size is None:
+                    raise ValueError(
+                        'The input size should be specified. '
+                        f'Please set the in_size attribute of the previous module: \n'
+                        f'{self.layers[-1]}'
+                    )
+                module = module(in_size=in_size)
+                assert isinstance(module, Module), 'The module should be an instance of Module.'
+                out_size = module.out_size
+            elif isinstance(module, ElementWiseBlock):
+                out_size = in_size
+            elif isinstance(module, Module):
+                out_size = module.out_size
+            elif callable(module):
+                out_size = in_size
+            else:
+                raise TypeError(f"Unsupported type {type(module)}. ")
+        except Exception as e:
+            raise BrainStateError(
+                f'Failed to format the module: \n'
+                f'{module}\n'
+                f'with input size: {in_size}\n'
+            ) from e
         return module, out_size

brainstate/nn/_normalizations.py

@@ -77,16 +77,16 @@ def weight_standardization(
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Standardize a weight matrix
         >>> w = jnp.ones((3, 4))
-        >>> w_std = bst.nn.weight_standardization(w)
+        >>> w_std = brainstate.nn.weight_standardization(w)
         >>>
         >>> # With custom gain
         >>> gain = jnp.ones((4,))
-        >>> w_std = bst.nn.weight_standardization(w, gain=gain)
+        >>> w_std = brainstate.nn.weight_standardization(w, gain=gain)
     """
     w = u.maybe_custom_array(w)
     if out_axis < 0:
@@ -551,11 +551,11 @@ class BatchNorm0d(_BatchNorm):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a BatchNorm0d layer
-        >>> layer = bst.nn.BatchNorm0d(in_size=(10,))
+        >>> layer = brainstate.nn.BatchNorm0d(in_size=(10,))
         >>>
         >>> # Apply normalization to a batch of data
         >>> x = jnp.ones((32, 10))  # batch_size=32, features=10
@@ -623,11 +623,11 @@ class BatchNorm1d(_BatchNorm):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a BatchNorm1d layer for sequence data
-        >>> layer = bst.nn.BatchNorm1d(in_size=(100, 64))  # length=100, channels=64
+        >>> layer = brainstate.nn.BatchNorm1d(in_size=(100, 64))  # length=100, channels=64
         >>>
         >>> # Apply normalization
         >>> x = jnp.ones((8, 100, 64))  # batch_size=8
@@ -693,11 +693,11 @@ class BatchNorm2d(_BatchNorm):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a BatchNorm2d layer for image data
-        >>> layer = bst.nn.BatchNorm2d(in_size=(28, 28, 3))  # 28x28 RGB images
+        >>> layer = brainstate.nn.BatchNorm2d(in_size=(28, 28, 3))  # 28x28 RGB images
         >>>
         >>> # Apply normalization
         >>> x = jnp.ones((16, 28, 28, 3))  # batch_size=16
@@ -763,11 +763,11 @@ class BatchNorm3d(_BatchNorm):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
         >>>
         >>> # Create a BatchNorm3d layer for volumetric data
-        >>> layer = bst.nn.BatchNorm3d(in_size=(32, 32, 32, 1))  # 32x32x32 volumes
+        >>> layer = brainstate.nn.BatchNorm3d(in_size=(32, 32, 32, 1))  # 32x32x32 volumes
         >>>
         >>> # Apply normalization
         >>> x = jnp.ones((4, 32, 32, 32, 1))  # batch_size=4
@@ -841,11 +841,11 @@ class LayerNorm(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>>
         >>> # Create a LayerNorm layer
-        >>> x = bst.random.normal(size=(3, 4, 5, 6))
-        >>> layer = bst.nn.LayerNorm(x.shape)
+        >>> x = brainstate.random.normal(size=(3, 4, 5, 6))
+        >>> layer = brainstate.nn.LayerNorm(x.shape)
         >>>
         >>> # Apply normalization
         >>> y = layer(x)
@@ -853,8 +853,8 @@ class LayerNorm(Module):
         (3, 4, 5, 6)
         >>>
         >>> # Normalize only the last dimension
-        >>> layer = bst.nn.LayerNorm((10, 20), reduction_axes=-1, feature_axes=-1)
-        >>> x = bst.random.normal((5, 10, 20))
+        >>> layer = brainstate.nn.LayerNorm((10, 20), reduction_axes=-1, feature_axes=-1)
+        >>> x = brainstate.random.normal((5, 10, 20))
         >>> y = layer(x)
     """
 
@@ -1005,11 +1005,11 @@ class RMSNorm(Module):
     --------
     .. code-block:: python
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>>
         >>> # Create an RMSNorm layer
-        >>> x = bst.random.normal(size=(5, 6))
-        >>> layer = bst.nn.RMSNorm(in_size=(6,))
+        >>> x = brainstate.random.normal(size=(5, 6))
+        >>> layer = brainstate.nn.RMSNorm(in_size=(6,))
         >>>
         >>> # Apply normalization
         >>> y = layer(x)
@@ -1017,8 +1017,8 @@ class RMSNorm(Module):
         (5, 6)
         >>>
         >>> # Without scaling
-        >>> layer = bst.nn.RMSNorm(in_size=(10,), use_scale=False)
-        >>> x = bst.random.normal((3, 10))
+        >>> layer = brainstate.nn.RMSNorm(in_size=(10,), use_scale=False)
+        >>> x = brainstate.random.normal((3, 10))
         >>> y = layer(x)
     """
 
@@ -1179,20 +1179,20 @@ class GroupNorm(Module):
     .. code-block:: python
 
         >>> import numpy as np
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>>
         >>> # Create a GroupNorm layer with 3 groups
-        >>> 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)
         >>> y = layer(x)
         >>>
         >>> # GroupNorm with num_groups=1 is equivalent to LayerNorm
-        >>> y1 = bst.nn.GroupNorm(x.shape, num_groups=1)(x)
-        >>> y2 = bst.nn.LayerNorm(x.shape, reduction_axes=(1, 2, 3))(x)
+        >>> y1 = 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(y1, y2, rtol=1e-5)
         >>>
         >>> # Specify group_size instead of num_groups
-        >>> layer = bst.nn.GroupNorm((12,), num_groups=None, group_size=4)
+        >>> layer = brainstate.nn.GroupNorm((12,), num_groups=None, group_size=4)
     """
 
     def __init__(

brainstate/random/__init__.py

@@ -259,12 +259,12 @@ References
 
 """
 
-from ._rand_funs import *
-from ._rand_funs import __all__ as __all_random__
-from ._rand_seed import *
-from ._rand_seed import __all__ as __all_seed__
-from ._rand_state import *
-from ._rand_state import __all__ as __all_state__
+from ._fun import *
+from ._fun import __all__ as __all_random__
+from ._seed import *
+from ._seed import __all__ as __all_seed__
+from ._state import *
+from ._state import __all__ as __all_state__
 
 __all__ = __all_random__ + __all_state__ + __all_seed__
 del __all_random__, __all_state__, __all_seed__

brainstate/random/{_rand_funs.py → _fun.py}

@@ -21,7 +21,7 @@ from typing import Optional
 import numpy as np
 
 from brainstate.typing import DTypeLike, Size, SeedOrKey
-from ._rand_state import RandomState, DEFAULT
+from ._state import RandomState, DEFAULT
 
 __all__ = [
     # numpy compatibility

brainstate/random/{_rand_funs_test.py → _fun_test.py}

@@ -437,13 +437,11 @@ class TestRandom(unittest.TestCase):
         a = brainstate.random.hypergeometric(10, 10, 10, 20)
         self.assertTupleEqual(a.shape, (20,))
 
-    @pytest.mark.skipif(platform.system() == 'Windows', reason='Windows jaxlib error')
     def test_hypergeometric2(self):
         brainstate.random.seed()
         a = brainstate.random.hypergeometric(8, [10, 4], [[5, 2], [5, 5]])
         self.assertTupleEqual(a.shape, (2, 2))
 
-    @pytest.mark.skipif(platform.system() == 'Windows', reason='Windows jaxlib error')
     def test_hypergeometric3(self):
         brainstate.random.seed()
         a = brainstate.random.hypergeometric(8, [10, 4], [[5, 2], [5, 5]], size=(3, 2, 2))