brainstate 0.1.0.post20250212__py2.py3-none-any.whl → 0.1.0.post20250216__py2.py3-none-any.whl

brainstate/nn/_dyn_impl/_dynamics_neuron_test.py

@@ -17,11 +17,10 @@
 
 from __future__ import annotations
 
-import unittest
-
 import brainunit as u
 import jax
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 from brainstate.nn import IF, LIF, ALIF

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -17,9 +17,8 @@
 
 from __future__ import annotations
 
-from typing import Optional
-
 import brainunit as u
+from typing import Optional
 
 from brainstate import init, environ
 from brainstate._state import ShortTermState, HiddenState

brainstate/nn/_dyn_impl/_dynamics_synapse_test.py

@@ -15,11 +15,10 @@
 
 from __future__ import annotations
 
-import unittest
-
 import brainunit as u
 import jax.numpy as jnp
 import pytest
+import unittest
 
 import brainstate as bst
 from brainstate.nn import Expon, STP, STD

brainstate/nn/_dyn_impl/_inputs.py

@@ -14,11 +14,10 @@
 # ==============================================================================
 from __future__ import annotations
 
-from typing import Union, Optional, Sequence, Callable
-
 import brainunit as u
 import jax
 import numpy as np
+from typing import Union, Optional, Sequence, Callable
 
 from brainstate import environ, init, random
 from brainstate._state import ShortTermState

brainstate/nn/_dyn_impl/_rate_rnns.py

@@ -17,9 +17,8 @@
 
 from __future__ import annotations
 
-from typing import Callable, Union
-
 import jax.numpy as jnp
+from typing import Callable, Union
 
 from brainstate import random, init, functional
 from brainstate._state import HiddenState, ParamState

brainstate/nn/_dyn_impl/_rate_rnns_test.py

@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 

brainstate/nn/_dyn_impl/_readout.py

@@ -17,11 +17,10 @@
 
 from __future__ import annotations
 
-import numbers
-from typing import Callable
-
 import brainunit as u
 import jax
+import numbers
+from typing import Callable
 
 from brainstate import environ, init, surrogate
 from brainstate._state import HiddenState, ParamState

brainstate/nn/_dyn_impl/_readout_test.py

@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 

brainstate/nn/_dynamics/_dynamics_base.py

@@ -34,10 +34,9 @@ For handling the delays:
 """
 from __future__ import annotations
 
-from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING
-
 import brainunit as u
 import numpy as np
+from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING
 
 from brainstate import environ
 from brainstate._state import State
@@ -420,7 +419,7 @@ class Dynamics(Module):
         else:
             raise TypeError(f'The input {dyn} should be an instance of {Dynamics} or a delayed initializer.')
 
-    def __leaf_fn__(self, name, value):
+    def __pretty_repr_item__(self, name, value):
         if name in ['_in_size', '_out_size', '_name', '_mode',
                     '_before_updates', '_after_updates', '_current_inputs', '_delta_inputs']:
             return (name, value) if value is None else (name[1:], value)  # skip the first `_`

brainstate/nn/_dynamics/_dynamics_base_test.py

@@ -17,9 +17,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import numpy as np
+import unittest
 
 import brainstate as bst
 

brainstate/nn/_dynamics/_state_delay.py

@@ -16,14 +16,14 @@
 from __future__ import annotations
 
 import math
-import numbers
-from functools import partial
-from typing import Optional, Dict, Callable, Union, Sequence
 
 import brainunit as u
 import jax
 import jax.numpy as jnp
+import numbers
 import numpy as np
+from functools import partial
+from typing import Optional, Dict, Callable, Union, Sequence
 
 from brainstate import environ
 from brainstate._state import ShortTermState, State

brainstate/nn/_dynamics/_synouts_test.py

@@ -15,11 +15,10 @@
 
 from __future__ import annotations
 
-import unittest
-
 import brainunit as u
 import jax.numpy as jnp
 import numpy as np
+import unittest
 
 import brainstate as bst
 

brainstate/nn/_elementwise/_dropout.py

@@ -16,11 +16,10 @@
 
 from __future__ import annotations
 
-from functools import partial
-from typing import Optional, Sequence
-
 import brainunit as u
 import jax.numpy as jnp
+from functools import partial
+from typing import Optional, Sequence
 
 from brainstate import random, environ, init
 from brainstate._state import ShortTermState
@@ -72,11 +71,11 @@ class Dropout(ElementWiseBlock):
             for dim in self.broadcast_dims:
                 broadcast_shape[dim] = 1
             keep_mask = random.bernoulli(self.prob, broadcast_shape)
-            keep_mask = jnp.broadcast_to(keep_mask, x.shape)
-            return jnp.where(
+            keep_mask = u.math.broadcast_to(keep_mask, x.shape)
+            return u.math.where(
                 keep_mask,
-                jnp.asarray(x / self.prob, dtype=dtype),
-                jnp.asarray(0., dtype=dtype)
+                u.math.asarray(x / self.prob, dtype=dtype),
+                u.math.asarray(0., dtype=dtype)
             )
         else:
             return x

brainstate/nn/_elementwise/_dropout_test.py

@@ -14,9 +14,8 @@
 # ==============================================================================
 
 
-import unittest
-
 import numpy as np
+import unittest
 
 import brainstate as bst
 

brainstate/nn/_elementwise/_elementwise.py

@@ -17,11 +17,10 @@
 
 from __future__ import annotations
 
-from typing import Optional
-
 import brainunit as u
 import jax.numpy as jnp
 import jax.typing
+from typing import Optional
 
 from brainstate import random, functional as F
 from brainstate._state import ParamState

brainstate/nn/_exp_euler.py

@@ -16,10 +16,9 @@
 
 from __future__ import annotations
 
-from typing import Callable
-
 import brainunit as u
 import jax.numpy as jnp
+from typing import Callable
 
 from brainstate import environ, random
 from brainstate.augment import vector_grad

brainstate/nn/_exp_euler_test.py

@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import brainunit as u
+import unittest
 
 import brainstate as bst
 

brainstate/nn/_interaction/_conv.py

@@ -18,10 +18,9 @@
 from __future__ import annotations
 
 import collections.abc
-from typing import Callable, Tuple, Union, Sequence, Optional, TypeVar
-
 import jax
 import jax.numpy as jnp
+from typing import Callable, Tuple, Union, Sequence, Optional, TypeVar
 
 from brainstate import init, functional
 from brainstate._state import ParamState

brainstate/nn/_interaction/_conv_test.py

@@ -235,5 +235,6 @@ class TestConvTranspose3d(parameterized.TestCase):
         y = conv_transpose_module(x)
         print(y.shape)
 
+
 if __name__ == '__main__':
     absltest.main()

brainstate/nn/_interaction/_linear.py

@@ -17,10 +17,9 @@
 
 from __future__ import annotations
 
-from typing import Callable, Union, Optional
-
 import brainunit as u
 import jax.numpy as jnp
+from typing import Callable, Union, Optional
 
 from brainstate import init, functional
 from brainstate._state import ParamState

brainstate/nn/_interaction/_linear_test.py

@@ -16,9 +16,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import brainunit as u
+import unittest
 from absl.testing import parameterized
 
 import brainstate as bst

brainstate/nn/_interaction/_normalizations.py

@@ -17,10 +17,9 @@
 
 from __future__ import annotations
 
-from typing import Callable, Union, Sequence, Optional, Any, Dict
-
 import jax
 import jax.numpy as jnp
+from typing import Callable, Union, Sequence, Optional, Any
 
 from brainstate import environ, init
 from brainstate._state import ParamState, BatchState

brainstate/nn/_interaction/_poolings.py

@@ -17,14 +17,13 @@
 
 from __future__ import annotations
 
-import functools
-from typing import Sequence, Optional
-from typing import Union, Tuple, Callable, List
-
 import brainunit as u
+import functools
 import jax
 import jax.numpy as jnp
 import numpy as np
+from typing import Sequence, Optional
+from typing import Union, Tuple, Callable, List
 
 from brainstate import environ
 from brainstate.nn._module import Module

brainstate/nn/_module.py

@@ -27,16 +27,15 @@ The basic classes include:
 """
 from __future__ import annotations
 
-import warnings
-from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING
-
 import numpy as np
+import warnings
+from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING, Callable
 
 from brainstate._state import State
 from brainstate.graph import Node, states, nodes, flatten
 from brainstate.mixin import ParamDescriber, ParamDesc
 from brainstate.typing import PathParts
-from brainstate.util import FlattedDict, NestedDict
+from brainstate.util import FlattedDict, NestedDict, BrainStateError
 
 # maximum integer
 max_int = np.iinfo(np.int32).max
@@ -226,7 +225,7 @@ class Module(Node, ParamDesc):
         """
         pass
 
-    def __leaf_fn__(self, name, value):
+    def __pretty_repr_item__(self, name, value):
         if name in ['_in_size', '_out_size', '_name']:
             return (name, value) if value is None else (name[1:], value)  # skip the first `_`
         return name, value
@@ -288,7 +287,7 @@ class Sequential(Module):
         in_size = first.out_size
         self.layers.append(first)
         for module in layers:
-            module, in_size = _format_module(module, in_size)
+            module, in_size = self._format_module(module, in_size)
             self.layers.append(module)
 
         # the input and output shape
@@ -301,7 +300,14 @@ class Sequential(Module):
         """Update function of a sequential model.
         """
         for m in self.layers:
-            x = m(x)
+            try:
+                x = m(x)
+            except Exception as e:
+                raise BrainStateError(
+                    f'The module \n'
+                    f'{m}\n'
+                    f'failed to update with input {x}\n'
+                ) from e
         return x
 
     def __getitem__(self, key: Union[int, slice]):
@@ -314,16 +320,54 @@ class Sequential(Module):
         else:
             raise KeyError(f'Unknown type of key: {type(key)}')
 
+    def append(self, layer: Callable):
+        """
+        Append a layer to the sequential model.
+
+        This method adds a new layer to the end of the sequential model. The layer can be
+        either a Module instance, an ElementWiseBlock instance, or a callable function. If the
+        layer is a callable function, it will be wrapped in an ElementWiseBlock instance.
+
+        Parameters:
+        ----------
+        layer : Callable
+            The layer to be appended to the sequential model. It can be a Module instance,
+            an ElementWiseBlock instance, or a callable function.
+
+        Raises:
+        -------
+        ValueError
+            If the sequential model is empty and the first layer is a callable function.
 
-def _format_module(module, in_size):
-    if isinstance(module, ParamDescriber):
-        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
-    else:
-        raise TypeError(f"Unsupported type {type(module)}. ")
-    return module, out_size
+        Returns:
+        --------
+        None
+            The method does not return any value. It modifies the sequential model by adding
+            the new layer to the end.
+        """
+        if len(self.layers) == 0:
+            raise ValueError('The first layer should be a module, not a function.')
+        module, in_size = self._format_module(layer, self.out_size)
+        self.layers.append(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)}. ")
+        return module, out_size

brainstate/nn/_module_test.py

@@ -15,10 +15,9 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax.numpy as jnp
 import jaxlib.xla_extension
+import unittest
 
 import brainstate as bst
 

brainstate/nn/metrics.py

@@ -16,13 +16,12 @@
 
 from __future__ import annotations
 
-import typing as tp
-from dataclasses import dataclass
-from functools import partial
-
 import jax
 import jax.numpy as jnp
 import numpy as np
+import typing as tp
+from dataclasses import dataclass
+from functools import partial
 
 from brainstate._state import State
 

brainstate/optim/_lr_scheduler.py

@@ -16,11 +16,10 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from typing import Sequence, Union
-
 import jax
 import jax.numpy as jnp
 import numpy as np
+from typing import Sequence, Union
 
 from brainstate import environ
 from brainstate._state import State, LongTermState

brainstate/optim/_lr_scheduler_test.py

@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 
@@ -37,7 +36,7 @@ class TestMultiStepLR(unittest.TestCase):
                 self.assertTrue(jnp.allclose(r, 0.0001))
 
     def test2(self):
-        lr = bst.transform.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
+        lr = bst.compile.jit(bst.optim.MultiStepLR(0.1, [10, 20, 30], gamma=0.1))
         for i in range(40):
             r = lr(i)
             if i < 10:

brainstate/optim/_optax_optimizer_test.py

@@ -15,10 +15,9 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax
 import optax
+import unittest
 
 import brainstate as bst
 

brainstate/optim/_sgd_optimizer.py

@@ -16,12 +16,11 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-import functools
-from typing import Union, Dict, Optional, Tuple, Any, TypeVar
-
 import brainunit as u
+import functools
 import jax
 import jax.numpy as jnp
+from typing import Union, Dict, Optional, Tuple, Any, TypeVar
 
 from brainstate import environ
 from brainstate._state import State, LongTermState, StateDictManager

brainstate/random/_rand_funs.py

@@ -17,9 +17,8 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from typing import Optional
-
 import numpy as np
+from typing import Optional
 
 from brainstate.typing import DTypeLike, Size, SeedOrKey
 from ._rand_state import RandomState, DEFAULT

brainstate/random/_rand_funs_test.py

@@ -15,13 +15,12 @@
 
 from __future__ import annotations
 
-import platform
-import unittest
-
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
+import platform
 import pytest
+import unittest
 
 import brainstate as bst
 

brainstate/random/_rand_seed.py

@@ -14,11 +14,10 @@
 # ==============================================================================
 from __future__ import annotations
 
-from contextlib import contextmanager
-from typing import Optional
-
 import jax
 import numpy as np
+from contextlib import contextmanager
+from typing import Optional
 
 from brainstate.typing import SeedOrKey
 from ._rand_state import RandomState, DEFAULT, use_prng_key

brainstate/random/_rand_seed_test.py

@@ -15,10 +15,9 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax.numpy as jnp
 import jax.random
+import unittest
 
 import brainstate as bst
 

brainstate/random/_rand_state.py

@@ -16,17 +16,16 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from functools import partial
-from operator import index
-from typing import Optional
-
 import brainunit as u
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 import numpy as np
+from functools import partial
 from jax import jit, vmap
 from jax import lax, core, dtypes
+from operator import index
+from typing import Optional
 
 from brainstate import environ
 from brainstate._state import State

brainstate/surrogate.py

@@ -21,7 +21,7 @@ import jax.numpy as jnp
 import jax.scipy as sci
 from jax.interpreters import batching, ad, mlir
 
-from brainstate.util._dict import PrettyObject
+from brainstate.util._pretty_pytree import PrettyObject
 
 if jax.__version_info__ < (0, 4, 38):
     from jax.core import Primitive
@@ -79,7 +79,10 @@ def _heaviside_imp(x, dx):
 
 
 def _heaviside_batching(args, axes):
-    return heaviside_p.bind(*args), tuple(axes)
+    x, dx = args
+    if axes[0] != axes[1]:
+        dx = batching.moveaxis(dx, axes[1], axes[0])
+    return heaviside_p.bind(x, dx), tuple([axes[0]])
 
 
 def _heaviside_jvp(primals, tangents):

brainstate/transform.py

@@ -15,6 +15,3 @@
 
 # alias for compilation and augmentation functions
 
-from .compile import *
-from .augment import *
-