brainstate 0.1.0.post20250212__py2.py3-none-any.whl → 0.1.0.post20250217__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

@@ -157,6 +157,11 @@ class Dynamics(Module):
         # 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']:
+            return None if value is None else (name[1:], value)  # skip the first `_`
+        return super().__pretty_repr_item__(name, value)
+
     @property
     def varshape(self):
         """The shape of variables in the neuron group."""
@@ -420,7 +425,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

@@ -28,7 +28,7 @@ The basic classes include:
 from __future__ import annotations
 
 import warnings
-from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING
+from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING, Callable
 
 import numpy as np
 
@@ -36,7 +36,7 @@ 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
@@ -113,7 +113,11 @@ class Module(Node, ParamDesc):
         """
         The function to specify the updating rule.
         """
-        raise NotImplementedError(f'Subclass of {self.__class__.__name__} must implement "update" function.')
+        raise NotImplementedError(
+            f'Subclass of {self.__class__.__name__} must implement "update" function. \n'
+            f'This instance is: \n'
+            f'{self}'
+        )
 
     def __call__(self, *args, **kwargs):
         return self.update(*args, **kwargs)
@@ -226,9 +230,9 @@ 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 None if value is None else (name[1:], value)  # skip the first `_`
         return name, value
 
 
@@ -288,7 +292,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 +305,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 +325,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.
 
-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
+        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.
+
+        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/_utils.py

@@ -0,0 +1,89 @@
+# Copyright 2025 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from typing import Union, Tuple
+
+from brainstate._state import ParamState
+from brainstate.util import PrettyTable
+from ._module import Module
+
+__all__ = [
+    "count_parameters",
+]
+
+
+def _format_parameter_count(num_params, precision=2):
+    if num_params < 1000:
+        return str(num_params)
+
+    suffixes = ['', 'K', 'M', 'B', 'T', 'P', 'E']
+    magnitude = 0
+    while abs(num_params) >= 1000:
+        magnitude += 1
+        num_params /= 1000.0
+
+    format_string = '{:.' + str(precision) + 'f}{}'
+    formatted_value = format_string.format(num_params, suffixes[magnitude])
+
+    # 检查是否接近 1000，如果是，尝试使用更大的基数
+    if magnitude < len(suffixes) - 1 and num_params >= 1000 * (1 - 10 ** (-precision)):
+        magnitude += 1
+        num_params /= 1000.0
+        formatted_value = format_string.format(num_params, suffixes[magnitude])
+
+    return formatted_value
+
+
+def count_parameters(
+    module: Module,
+    precision: int = 2,
+    return_table: bool = False,
+) -> Union[Tuple[PrettyTable, int], int]:
+    """
+    Count and display the number of trainable parameters in a neural network model.
+
+    This function iterates through all the parameters of the given model,
+    counts the number of parameters for each module, and displays them in a table.
+    It also calculates and returns the total number of trainable parameters.
+
+    Parameters:
+    -----------
+    model : bst.nn.Module
+        The neural network model for which to count parameters.
+
+    Returns:
+    --------
+    int
+        The total number of trainable parameters in the model.
+
+    Prints:
+    -------
+    A pretty-formatted table showing the number of parameters for each module,
+    followed by the total number of trainable parameters.
+    """
+    assert isinstance(module, Module), "Input must be a neural network module"  # noqa: E501
+    table = PrettyTable(["Modules", "Parameters"])
+    total_params = 0
+    for name, parameter in module.states(ParamState).items():
+        param = parameter.numel()
+        table.add_row([name, _format_parameter_count(param, precision=precision)])
+        total_params += param
+    table.add_row(["Total", _format_parameter_count(total_params, precision=precision)])
+    print(table)
+    if return_table:
+        return table, total_params
+    return total_params

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