brainstate 0.0.2.post20240826__py2.py3-none-any.whl → 0.0.2.post20240910__py2.py3-none-any.whl

brainstate/_module.py

@@ -1597,6 +1597,6 @@ def _get_delay(delay_time, delay_step):
       delay_time = delay_step * environ.get_dt()
   else:
     assert delay_step is None, '"delay_step" should be None if "delay_time" is given.'
-    assert isinstance(delay_time, (int, float))
+    # assert isinstance(delay_time, (int, float))
     delay_step = math.ceil(delay_time / environ.get_dt())
   return delay_time, delay_step

brainstate/_state.py

@@ -234,6 +234,8 @@ class ShapeDtype:
   def __init__(self, shape, dtype):
     self.shape = shape
     self.dtype = dtype
+    self.ndim = len(shape)
+    self.size = np.prod(shape)
 
   def __repr__(self):
     return f'{self.dtype}{list(self.shape)}'

brainstate/nn/_projection/_align_post.py

@@ -15,6 +15,8 @@
 
 from typing import Optional, Union
 
+
+import brainunit as u
 from brainstate._module import (register_delay_of_target,
                                 Projection,
                                 Module,
@@ -278,11 +280,19 @@ class FullProjAlignPostMg(Projection):
     self.comm = comm
 
     # delay initialization
-    if delay is not None and delay > 0.:
-      delay_cls = register_delay_of_target(pre)
-      delay_cls.register_entry(self.name, delay)
-      self.delay = delay_cls
-      self.has_delay = True
+    if delay is not None:
+      if isinstance(delay, u.Quantity):
+        has_delay = delay.mantissa > 0.
+      else:
+        has_delay = delay > 0.
+      if has_delay:
+        delay_cls = register_delay_of_target(pre)
+        delay_cls.register_entry(self.name, delay)
+        self.delay = delay_cls
+        self.has_delay = True
+      else:
+        self.delay = None
+        self.has_delay = False
     else:
       self.delay = None
       self.has_delay = False
@@ -502,11 +512,19 @@ class FullProjAlignPost(Projection):
     self.out = out
 
     # delay initialization
-    if delay is not None and delay > 0.:
-      delay_cls = register_delay_of_target(pre)
-      delay_cls.register_entry(self.name, delay)
-      self.delay = delay_cls
-      self.has_delay = True
+    if delay is not None:
+      if isinstance(delay, u.Quantity):
+        has_delay = delay.mantissa > 0.
+      else:
+        has_delay = delay > 0.
+      if has_delay:
+        delay_cls = register_delay_of_target(pre)
+        delay_cls.register_entry(self.name, delay)
+        self.delay = delay_cls
+        self.has_delay = True
+      else:
+        self.delay = None
+        self.has_delay = False
     else:
       self.delay = None
       self.has_delay = False

brainstate/surrogate.py

@@ -158,6 +158,9 @@ class Sigmoid(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def sigmoid(
     x: jax.Array,
@@ -243,6 +246,9 @@ class PiecewiseQuadratic(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def piecewise_quadratic(
     x: jax.Array,
@@ -339,6 +345,9 @@ class PiecewiseExp(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def piecewise_exp(
     x: jax.Array,
@@ -426,6 +435,9 @@ class SoftSign(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def soft_sign(
     x: jax.Array,
@@ -508,6 +520,9 @@ class Arctan(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def arctan(
     x: jax.Array,
@@ -589,6 +604,9 @@ class NonzeroSignLog(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def nonzero_sign_log(
     x: jax.Array,
@@ -683,6 +701,9 @@ class ERF(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def erf(
     x: jax.Array,
@@ -780,6 +801,9 @@ class PiecewiseLeakyRelu(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(c={self.c}, w={self.w})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.c, self.w))
+
 
 def piecewise_leaky_relu(
     x: jax.Array,
@@ -898,6 +922,9 @@ class SquarewaveFourierSeries(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(n={self.n}, t_period={self.t_period})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.n, self.t_period))
+
 
 def squarewave_fourier_series(
     x: jax.Array,
@@ -988,6 +1015,9 @@ class S2NN(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta}, epsilon={self.epsilon})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.beta, self.epsilon))
+
 
 def s2nn(
     x: jax.Array,
@@ -1089,6 +1119,9 @@ class QPseudoSpike(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def q_pseudo_spike(
     x: jax.Array,
@@ -1178,6 +1211,9 @@ class LeakyRelu(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.beta))
+
 
 def leaky_relu(
     x: jax.Array,
@@ -1277,6 +1313,9 @@ class LogTailedRelu(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def log_tailed_relu(
     x: jax.Array,
@@ -1368,6 +1407,9 @@ class ReluGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, width={self.width})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.width))
+
 
 def relu_grad(
     x: jax.Array,
@@ -1446,6 +1488,9 @@ class GaussianGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha}, sigma={self.sigma})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha, self.sigma))
+
 
 def gaussian_grad(
     x: jax.Array,
@@ -1530,6 +1575,9 @@ class MultiGaussianGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(h={self.h}, s={self.s}, sigma={self.sigma}, scale={self.scale})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.h, self.s, self.sigma, self.scale))
+
 
 def multi_gaussian_grad(
     x: jax.Array,
@@ -1615,6 +1663,9 @@ class InvSquareGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def inv_square_grad(
     x: jax.Array,
@@ -1685,6 +1736,9 @@ class SlayerGrad(Surrogate):
   def __repr__(self):
     return f'{self.__class__.__name__}(alpha={self.alpha})'
 
+  def __hash__(self):
+    return hash((self.__class__, self.alpha))
+
 
 def slayer_grad(
     x: jax.Array,

brainstate/transform/_loop_collect_return.py

@@ -418,7 +418,7 @@ def checkpointed_for_loop(
     pbar: Optional[ProgressBar] = None,
 ):
   """
-  ``for-loop`` control flow with :py:class:`~.State` with a checkpointed version.
+  ``for-loop`` control flow with :py:class:`~.State` with a checkpointed version, similar to :py:func:`for_loop`.
 
   Args:
     f: a Python function to be scanned of type ``c -> a -> (c, b)``, meaning

{brainstate-0.0.2.post20240826.dist-info → brainstate-0.0.2.post20240910.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.2.post20240826
+Version: 0.0.2.post20240910
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BDP
@@ -78,16 +78,9 @@ pip install brainstate --upgrade
 The official documentation is hosted on Read the Docs: [https://brainstate.readthedocs.io/](https://brainstate.readthedocs.io/)
 
 
-## See also the BDP ecosystem
-
-- [``brainstate``](https://github.com/brainpy/brainstate): A ``State``-based transformation system for brain dynamics programming.
-
-- [``brainunit``](https://github.com/brainpy/brainunit): The unit system for brain dynamics programming.
 
-- [``braintaichi``](https://github.com/brainpy/braintaichi): Leveraging Taichi Lang to customize brain dynamics operators.
-
-- [``brainscale``](https://github.com/brainpy/brainscale): The scalable online learning framework for biological neural networks.
+## See also the BDP ecosystem
 
-- [``braintools``](https://github.com/brainpy/braintools): The toolbox for the brain dynamics simulation, training and analysis.
+We are building the BDP ecosystem: https://ecosystem-for-brain-dynamics.readthedocs.io/
 
 

{brainstate-0.0.2.post20240826.dist-info → brainstate-0.0.2.post20240910.dist-info}/RECORD

@@ -1,8 +1,8 @@
 brainstate/__init__.py,sha256=zipNSih9Tyvi4-5cXqNPGsDF7VeestkLp-lcjJ4-dA0,1408
-brainstate/_module.py,sha256=L0cwF_6p9cSZlWGi33Mb5s4vUnbfSeQi2TPUbvsGCzo,52461
+brainstate/_module.py,sha256=ULfItiqiQoIK1YUYfkasmyh8Rj4PoYJP7cxyuphEnIo,52463
 brainstate/_module_test.py,sha256=oQaoaZBTo1o3wHrMEJTInQCc7RdcVs1gcfQGvdSb1SI,7843
 brainstate/_random_for_unit.py,sha256=1rHr7gfH_bYrJfpxbDhQUk_j00Yosx-GzyZCXrLxsd0,2007
-brainstate/_state.py,sha256=C0widCOj_ca6zfqh95jzFXf_G5vi0hJyuQ5GIqEqOUs,12102
+brainstate/_state.py,sha256=J5j5NujvqU3Ftd_m_u_3mz4xWw81mdYgdzltrdJSy8o,12162
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
 brainstate/environ.py,sha256=k0p1oyi9jbsPfuvqrPL-_zgSd7VW3LRs0LboxlaaIfc,11806
@@ -10,7 +10,7 @@ brainstate/mixin.py,sha256=OumTTSVyYSbtudjfS_MRThsBaeVJ_0JggeMClY7xtBA,10758
 brainstate/mixin_test.py,sha256=-Ej9oUOu8O1M4oy37SVMj7xNRYhHHyAHwrjS_aISayo,2923
 brainstate/random.py,sha256=pTZvTH06hv08_TpwzAWCqAjy-8oNGmB6-Jp6MKfkLaY,188087
 brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,17866
-brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
+brainstate/surrogate.py,sha256=6AO79JOOs-X5x0FT0EDqO9lNtjJZAs26H4mljgpTvAw,45197
 brainstate/typing.py,sha256=szCYee9R15YQfsEAQOx95_LqfrD9AYuE5dfTBTPd8sg,9165
 brainstate/util.py,sha256=y-6eX1z3EMyg6pfZt4YdDalOnJ3HDAT1IPBCJDp-gQI,19876
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
@@ -38,7 +38,7 @@ brainstate/nn/_rate_rnns.py,sha256=Cebhy57UWzfwrCfq0v2qLDegmb__mXL5ht750y4aTro,1
 brainstate/nn/_readout.py,sha256=jsQwhVnrJICKw4wFq-Du2AORPb_XXz_tZ4cURcckU-E,4240
 brainstate/nn/_synouts.py,sha256=gi3EyKlzt4UoyghwvNIr03r7YabZyl1idbq9aYG8zYM,4379
 brainstate/nn/_projection/__init__.py,sha256=L6svNHTb8BDh2rdX2eYmcx_NdscSdKykkQbzpdCSkTA,1207
-brainstate/nn/_projection/_align_post.py,sha256=tV_-nizkiRy2B0bBsB4UFUANFUe1IjFZ8hDM9TGpSaU,20655
+brainstate/nn/_projection/_align_post.py,sha256=S1huNBq3NkOfwrr7SXgTU6JvQk7KPVv86XwJ5iyvaBI,21106
 brainstate/nn/_projection/_align_pre.py,sha256=_wjdj8muuv2_fSW9m3KBUVjNkBg28BUmz3qZ9IA1rUM,24597
 brainstate/nn/_projection/_delta.py,sha256=KT8ySo3n_Q_7swzOH-ISDf0x9rjMkiv99H-vqeQZDR8,7122
 brainstate/nn/_projection/_utils.py,sha256=UcmELOqsINgqJr7eC5BSNNteyZ--1lyGjhUTJfxyMmA,813
@@ -56,15 +56,15 @@ brainstate/transform/_error_if.py,sha256=0JThfFqt9B3K3H6mS84qecBS22yTi3-FPzviaYa
 brainstate/transform/_error_if_test.py,sha256=kQZujlgr9bYnL-Vf7x4Zfc7jJk7rCLNVu-bsiry40dQ,1874
 brainstate/transform/_jit.py,sha256=sjQHFV8Tt75fpdl12jjPRDPT92_IZxBBJAG4gapdbNQ,11471
 brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJRLb4,2868
-brainstate/transform/_loop_collect_return.py,sha256=aUZSK5MX4stVP8Te4R8glm2SdP18rUUfHjcV4TXOPC8,20768
+brainstate/transform/_loop_collect_return.py,sha256=8X6-3T3YoL_Buph9LiGASdrqPnRhsgsH9GQg1wcRos0,20800
 brainstate/transform/_loop_no_collection.py,sha256=p2vHoNNesDH2cM7b5LgLzSv90M8iNQPkRZEl0jhf7yA,6476
 brainstate/transform/_make_jaxpr.py,sha256=ZkrOZu4_0xcILuPUA3RFEkorJ-xbDuDtXorJI_qVThE,30450
 brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
 brainstate/transform/_mapping.py,sha256=G9XUsD1xKLCprwwE0wv0gSXS0NYZ-ZIsv-PKKRlOoTA,3821
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
 brainstate/transform/_unvmap.py,sha256=8Se_23QrwDdcJpFcUnnMgD6EP-4XylbhP9K5TDhW358,3311
-brainstate-0.0.2.post20240826.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.2.post20240826.dist-info/METADATA,sha256=unkbYHiPHAtNHGENoSt47mA0N3cXWsCafRC8Fo2NPyk,3851
-brainstate-0.0.2.post20240826.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.2.post20240826.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.2.post20240826.dist-info/RECORD,,
+brainstate-0.0.2.post20240910.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.post20240910.dist-info/METADATA,sha256=gAdKRqW3BdiBR-xvAOdqwKiex6eXywKX0bieOW2_ZZQ,3311
+brainstate-0.0.2.post20240910.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.2.post20240910.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.post20240910.dist-info/RECORD,,