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

brainstate/nn/__init__.py

@@ -31,6 +31,8 @@ from ._interaction import *
 from ._interaction import __all__ as interaction_all
 from ._module import *
 from ._module import __all__ as module_all
+from ._utils import *
+from ._utils import __all__ as utils_all
 
 __all__ = (
     ['metrics']
@@ -42,6 +44,7 @@ __all__ = (
     + module_all
     + exp_euler_all
     + interaction_all
+    + utils_all
 )
 
 del (
@@ -53,4 +56,5 @@ del (
     module_all,
     exp_euler_all,
     interaction_all,
+    utils_all,
 )

brainstate/nn/_common.py

@@ -174,7 +174,6 @@ class Vmap(Module):
             axis_size=axis_size,
         )
         def vmap_run(*args, **kwargs):
-            vmap_states_ = vmap_states
             return module(*args, **kwargs)
 
         # vmapped module

brainstate/nn/_dynamics/_dynamics_base.py

@@ -34,9 +34,10 @@ 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
@@ -156,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."""

brainstate/nn/_module.py

@@ -27,10 +27,11 @@ The basic classes include:
 """
 from __future__ import annotations
 
-import numpy as np
 import warnings
 from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING, Callable
 
+import numpy as np
+
 from brainstate._state import State
 from brainstate.graph import Node, states, nodes, flatten
 from brainstate.mixin import ParamDescriber, ParamDesc
@@ -112,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)
@@ -227,7 +232,7 @@ class Module(Node, ParamDesc):
 
     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
 
 

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/surrogate.py

@@ -132,9 +132,6 @@ class Surrogate(PrettyObject):
         dx = self.surrogate_grad(x)
         return heaviside_p.bind(x, dx)[0]
 
-    def __repr__(self):
-        return f'{self.__class__.__name__}()'
-
     def surrogate_fun(self, x) -> jax.Array:
         """The surrogate function."""
         raise NotImplementedError

{brainstate-0.1.0.post20250216.dist-info → brainstate-0.1.0.post20250217.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250216
+Version: 0.1.0.post20250217
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.0.post20250216.dist-info → brainstate-0.1.0.post20250217.dist-info}/RECORD

@@ -6,7 +6,7 @@ brainstate/environ.py,sha256=PllYYZKqany3G7NzIwoUPplLAePbyza6kJGXTPgJK-c,17698
 brainstate/environ_test.py,sha256=khZ_-SUJL6rQCgndeYV98ruUIHGTwFDtITrOs_olmuo,2043
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
-brainstate/surrogate.py,sha256=OiI7l51s-yjxOdPjTAwWyBljDq9hRTadjZGpCF8zVkc,53654
+brainstate/surrogate.py,sha256=wWYw-TxaFxHVneXuHjWD1UtTcOTk3XRSnhRtUkt_Hb8,53580
 brainstate/transform.py,sha256=vZWzO4F7qTsXL_SiVQPlTz0l9b_hRo9D-igETfgCTy0,758
 brainstate/typing.py,sha256=988gX1tvwtyYnYjmej90OaRxoMoBIPO0-DSrXXGxojM,10523
 brainstate/augment/__init__.py,sha256=Q9-JIwQ1FNn8VLS1MA9MrSylbvUjWSw98whrI3NIuKo,1229
@@ -53,14 +53,15 @@ brainstate/init/_random_inits.py,sha256=gyy9481ju7VEi-SFbSRU5iBACaHnf4wjI0596FNu
 brainstate/init/_random_inits_test.py,sha256=lBL2RQdBSZ88Zqz4IMdbHJMvDi7ooZq6caCpHfNtIRk,5197
 brainstate/init/_regular_inits.py,sha256=DmVMajugfyYFNUMzgFdDKMvbBu9hMWxkfDd-50uhoLg,3187
 brainstate/init/_regular_inits_test.py,sha256=tJl4aOkclllJIfKzJTbc0cfYCw2SoBsx8_G123RnqbU,1842
-brainstate/nn/__init__.py,sha256=Eb9q19tR29RAAvpjAQWubxfubzMzCY2gAHVpi6zSBjI,1725
+brainstate/nn/__init__.py,sha256=ar1hDUYbSO6oadMpbuS9FWZvZB_iyFzM8CwMK-RNDzM,1823
 brainstate/nn/_collective_ops.py,sha256=yQNBnh-XVEFnTg-Ga14mHOCGtGxiTkL9MYKdNjJF1BI,17535
 brainstate/nn/_collective_ops_test.py,sha256=yW7NNYsGFglFRFkqVlpGSY6WLnU-h8GlK6wCmG5jtRc,1189
-brainstate/nn/_common.py,sha256=Decmt8uwhIEJ5ODpl8gwUXSjSTYQcP9n0i7p3kt7eIo,7176
+brainstate/nn/_common.py,sha256=XQw0i0sH3Y_qUwHSMC7G9VQnDj-RuuTh1Ul-xRIPxxc,7136
 brainstate/nn/_exp_euler.py,sha256=cRgPNcjMs2C9x_8JabtYz5hm_FwqbiJ_U1VfRHYIlrE,3519
 brainstate/nn/_exp_euler_test.py,sha256=kvPf009DMYtla2uedKVKrPTHDyMTBepjlfsk5vDHqhI,1240
-brainstate/nn/_module.py,sha256=wJvgwztwjIQT_1o4o-DSjT081nEpfLDfmGWvg_-K5Ro,12764
+brainstate/nn/_module.py,sha256=vrukVI0ylbymzilh9BZtb-d9dnsBsykqanUNTx9Eb6Y,12844
 brainstate/nn/_module_test.py,sha256=UrVA85fo0KVFN9ApPkxkRcvtXEskWOXPzZIBa4JSFo0,8891
+brainstate/nn/_utils.py,sha256=epfELIy1COgdS9z5be-fmbFhagNugcIHpw4ww-HlkSY,3123
 brainstate/nn/metrics.py,sha256=p7eVwd5y8r0N5rMws-zOS_KaZCLOMdrXyQvLnoJeq1w,14736
 brainstate/nn/_dyn_impl/__init__.py,sha256=Oazar7h89dp1WA2Vx4Tj7gCBhxJKH4LAUEABkBEG7vU,1462
 brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=mcDxVZlk56NAEkR6xcE74hOZ9up8Rua4SvKEeAhJKU4,10925
@@ -74,7 +75,7 @@ brainstate/nn/_dyn_impl/_rate_rnns_test.py,sha256=EYUajj50PL1V_yuQIANs6sTWfSDrXK
 brainstate/nn/_dyn_impl/_readout.py,sha256=iXqOahtDaLgMuMYXdMT-0scMPLHK1-fYeBb8SfEShY0,4368
 brainstate/nn/_dyn_impl/_readout_test.py,sha256=QYXoWlTXVwJoIVqAXm5UYF5bjHUMkY4bKqWWyXzXF10,2107
 brainstate/nn/_dynamics/__init__.py,sha256=j1HSWu01wf5-KjSaNhBC9utVGDALOhUsFPrLPcPPDsM,1208
-brainstate/nn/_dynamics/_dynamics_base.py,sha256=XADbPoPjNR_Uwx4oSwoA-qMuYwFR3gJRCnR4_gHtq4w,21994
+brainstate/nn/_dynamics/_dynamics_base.py,sha256=rLjTgA_826EfD1OZ-NoEoO11EBNJOpY_Fq2YfDKXRe4,22288
 brainstate/nn/_dynamics/_dynamics_base_test.py,sha256=Sk6sSqJK_yesI-6Fb_x7gqMsoc0-RUU9GsGZu2jVsxU,2719
 brainstate/nn/_dynamics/_projection_base.py,sha256=jYe3WdBMgz2TJkcxPWEkyK7OA4IR1ChISd2GTfM6U2o,13528
 brainstate/nn/_dynamics/_state_delay.py,sha256=qKF1YelGXeBlImhIBdZIC0CAg2dV0o_gkaxS-R1N3qE,16905
@@ -120,8 +121,8 @@ brainstate/util/_pretty_table.py,sha256=NM_6VAW6oL9jojsK0-RkQGHnDzLy_fn_hgzl5R8o
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=F5GfFURITAIYTwf17_xypkZU1wvoL4dUCviPnr_eCtw,17515
 brainstate/util/filter.py,sha256=Zw0H42NwAi2P7dBr3ISv2VpkB5jqoWnV4Kpd61gq66o,14126
-brainstate-0.1.0.post20250216.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250216.dist-info/METADATA,sha256=pLMfvR_kUK1DsX9tcUEvX1jCX2zdWj8YUx6YrFF7iug,3585
-brainstate-0.1.0.post20250216.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250216.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250216.dist-info/RECORD,,
+brainstate-0.1.0.post20250217.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250217.dist-info/METADATA,sha256=eBHppB4DxysakCjcinzG4vVSxI_eS08l7LkSepmGGXI,3585
+brainstate-0.1.0.post20250217.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250217.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250217.dist-info/RECORD,,