brainstate 0.0.1.1.post20240708__py2.py3-none-any.whl → 0.0.1.1.post20240802__py2.py3-none-any.whl

brainstate/_state.py

@@ -15,7 +15,7 @@
 
 import contextlib
 import threading
-from typing import Any, Tuple, Dict, List, Callable
+from typing import Any, Tuple, Dict, List, Callable, Optional
 
 import jax
 import numpy as np
@@ -108,9 +108,9 @@ class State(object):
     value: PyTree. It can be anything as a pyTree.
   """
   __module__ = 'brainstate'
-  __slots__ = ('_value', '_tree', '_level', '_source_info', '_check_tree')
+  __slots__ = ('_value', '_name', '_tree', '_level', '_source_info', '_check_tree')
 
-  def __init__(self, value: PyTree):
+  def __init__(self, value: PyTree, name: Optional[str] = None):
     if isinstance(value, State):
       value = value.value
     self._value = value
@@ -118,6 +118,21 @@ class State(object):
     self._check_tree = False
     self._level = len(thread_local_stack.stack)
     self._source_info = source_info_util.current()
+    self._name = name
+
+  @property
+  def name(self) -> Optional[str]:
+    """
+    The name of the state.
+    """
+    return self._name
+
+  @name.setter
+  def name(self, name: str) -> None:
+    """
+    Set the name of the state.
+    """
+    self._name = name
 
   @property
   def value(self) -> PyTree:
@@ -210,7 +225,10 @@ class State(object):
     leaves, tree = jax.tree.flatten(self._value)
     leaves_info = [ShapeDtype(leaf.shape, leaf.dtype) for leaf in leaves]
     tree_info = jax.tree.unflatten(tree, leaves_info)
-    return f'{self.__class__.__name__}({tree_info})'
+    if self.name is None:
+      return f'{self.__class__.__name__}({tree_info})'
+    else:
+      return f'{self.__class__.__name__}({self.name}: {tree_info})'
 
 
 class ShapeDtype:

brainstate/init/_generic.py

@@ -83,7 +83,7 @@ def _expand_params_to_match_sizes(params, sizes):
 
 
 def param(
-    parameter: Union[Callable, ArrayLike],
+    parameter: Union[Callable, ArrayLike, State],
     sizes: Union[int, Sequence[int]],
     batch_size: Optional[int] = None,
     allow_none: bool = True,

brainstate/mixin.py

@@ -207,7 +207,7 @@ class _JointGenericAlias(_UnionGenericAlias, _root=True):
 
 @_SpecialForm
 def JointTypes(self, parameters):
-  """All of types; AllOfTypes[X, Y] means both X and Y.
+  """Joint types; JointTypes[X, Y] means both X and Y.
 
   To define a union, use e.g. Union[int, str].
 
@@ -216,28 +216,28 @@ def JointTypes(self, parameters):
   - None as an argument is a special case and is replaced by `type(None)`.
   - Unions of unions are flattened, e.g.::
 
-      AllOfTypes[AllOfTypes[int, str], float] == AllOfTypes[int, str, float]
+      JointTypes[JointTypes[int, str], float] == JointTypes[int, str, float]
 
   - Unions of a single argument vanish, e.g.::
 
-      AllOfTypes[int] == int  # The constructor actually returns int
+      JointTypes[int] == int  # The constructor actually returns int
 
   - Redundant arguments are skipped, e.g.::
 
-      AllOfTypes[int, str, int] == AllOfTypes[int, str]
+      JointTypes[int, str, int] == JointTypes[int, str]
 
   - When comparing unions, the argument order is ignored, e.g.::
 
-      AllOfTypes[int, str] == AllOfTypes[str, int]
+      JointTypes[int, str] == JointTypes[str, int]
 
-  - You cannot subclass or instantiate a AllOfTypes.
-  - You can use Optional[X] as a shorthand for AllOfTypes[X, None].
+  - You cannot subclass or instantiate a JointTypes.
+  - You can use Optional[X] as a shorthand for JointTypes[X, None].
   """
   if parameters == ():
     raise TypeError("Cannot take a Joint of no types.")
   if not isinstance(parameters, tuple):
     parameters = (parameters,)
-  msg = "AllOfTypes[arg, ...]: each arg must be a type."
+  msg = "JointTypes[arg, ...]: each arg must be a type."
   parameters = tuple(_type_check(p, msg) for p in parameters)
   parameters = _remove_dups_flatten(parameters)
   if len(parameters) == 1:

brainstate/mixin_test.py

@@ -30,6 +30,8 @@ class TestMixin(unittest.TestCase):
     self.assertTrue(bc.mixin.Training)
 
 
+
+
 class TestMode(unittest.TestCase):
   def test_JointMode(self):
     a = bc.mixin.JointMode(bc.mixin.Batching(), bc.mixin.Training())

brainstate/nn/_elementwise.py

@@ -1139,13 +1139,13 @@ class Dropout(Module, ElementWiseBlock):
       name: Optional[str] = None
   ) -> None:
     super().__init__(mode=mode, name=name)
-    assert 0. <= prob < 1., f"Dropout probability must be in the range [0, 1). But got {prob}."
+    assert 0. <= prob <= 1., f"Dropout probability must be in the range [0, 1]. But got {prob}."
     self.prob = prob
 
   def __call__(self, x):
     dtype = bu.math.get_dtype(x)
     fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
-    if fit_phase:
+    if fit_phase and self.prob < 1.:
       keep_mask = random.bernoulli(self.prob, x.shape)
       return jnp.where(keep_mask,
                        jnp.asarray(x / self.prob, dtype=dtype),

brainstate/random.py

@@ -1167,23 +1167,32 @@ def default_rng(seed_or_key=None, clone: bool = True) -> RandomState:
     return RandomState(seed_or_key)
 
 
-def seed(seed: int = None):
+def seed(seed_or_key: int = None):
   """Sets a new random seed.
 
   Parameters
   ----------
-  seed: int, optional
-    The random seed.
+  seed_or_key: int, optional
+    The random seed (an integer) or jax random key.
   """
   with jax.ensure_compile_time_eval():
-    if seed is None:
-      seed = np.random.randint(0, 100000)
-    np.random.seed(seed)
-  DEFAULT.seed(seed)
+    if seed_or_key is None:
+      seed_or_key = np.random.randint(0, 100000)
+
+  # numpy random seed
+  if np.size(seed_or_key) == 1:  # seed
+    np.random.seed(seed_or_key)
+  elif np.size(seed_or_key) == 2:  # jax random key
+    np.random.seed(seed_or_key[0])
+  else:
+    raise ValueError(f"seed_or_key should be an integer or a tuple of two integers.")
+
+  # jax random seed
+  DEFAULT.seed(seed_or_key)
 
 
 @contextmanager
-def seed_context(seed: int):
+def seed_context(seed_or_key: SeedOrKey):
   """
   A context manager that sets the random seed for the duration of the block.
 
@@ -1206,16 +1215,19 @@ def seed_context(seed: int):
      The context manager does not only set the seed for the AX random state, but also for the numpy random state.
 
   Args:
-    seed: The seed (an integer).
+    seed_or_key: The seed (an integer) or jax random key.
 
-  Returns:
-    The random state.
   """
   old_jrand_key = DEFAULT.value
   old_np_state = np.random.get_state()
   try:
-    np.random.seed(seed)
-    DEFAULT.seed(seed)
+    if np.size(seed_or_key) == 1:  # seed
+      np.random.seed(seed_or_key)
+    elif np.size(seed_or_key) == 2:  # jax random key
+      np.random.seed(seed_or_key[0])
+    else:
+      raise ValueError(f"seed_or_key should be an integer or a tuple of two integers.")
+    DEFAULT.seed(seed_or_key)
     yield
   finally:
     np.random.set_state(old_np_state)
@@ -1223,7 +1235,8 @@ def seed_context(seed: int):
 
 
 def rand(*dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
-  r"""Random values in a given shape.
+  r"""
+  Random values in a given shape.
 
   .. note::
       This is a convenience function for users porting code from Matlab,

brainstate/transform/_control.py

@@ -25,7 +25,7 @@ import jax.numpy as jnp
 import numpy as np
 
 from brainstate._utils import set_module_as
-from ._jit_error import jit_error
+from ._jit_error import jit_error, remove_vmap
 from ._make_jaxpr import StatefulFunction, _assign_state_values
 from ._progress_bar import ProgressBar
 
@@ -347,7 +347,7 @@ def _wrap_fun_with_pbar(fun, pbar_runner):
   def new_fun(new_carry, inputs):
     i, old_carry = new_carry
     old_carry, old_outputs = fun(old_carry, inputs)
-    pbar_runner(i)
+    pbar_runner(remove_vmap(i, op='none'))
     return (i + 1, old_carry), old_outputs
 
   return new_fun

{brainstate-0.0.1.1.post20240708.dist-info → brainstate-0.0.1.1.post20240802.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.1.1.post20240708
+Version: 0.0.1.1.post20240802
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
-Author: BrainPy Team
+Author: BDP
 Author-email: BrainPy Team <chao.brain@qq.com>
 License: Apache-2.0 license
 Project-URL: homepage, http://github.com/brainpy

{brainstate-0.0.1.1.post20240708.dist-info → brainstate-0.0.1.1.post20240802.dist-info}/RECORD

@@ -2,13 +2,13 @@ brainstate/__init__.py,sha256=oxslZrm6wxtBQDqwJFb2BaAKZFmnp4d_esDkaeuGMWE,1410
 brainstate/_module.py,sha256=UjhfmY26VHQ-kF6U4l68AslGU6vhfD-dR7gF-4Io5ic,52520
 brainstate/_module_test.py,sha256=oQaoaZBTo1o3wHrMEJTInQCc7RdcVs1gcfQGvdSb1SI,7843
 brainstate/_random_for_unit.py,sha256=eW4NJkX27VCCNWUwAlyt2otkeEthGKOpUoX6XJ6i95Y,1946
-brainstate/_state.py,sha256=ykQluBkdKIcQsd7-pU_UlpnByWFof_8TYQl1hGn7HS8,11629
+brainstate/_state.py,sha256=t4lEikvxTfeL2TW0chLUvsQuuRoJSO-iXylUydl1i7k,12057
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
 brainstate/environ.py,sha256=LwRwnFaTbv8l7nHRIbSV46WzcN7pGLQFhT_xDUox2yA,10240
-brainstate/mixin.py,sha256=f0XovWlyTOQKC63WTtzSbztRSld3pKEHy9-I1gVZWLg,10748
-brainstate/mixin_test.py,sha256=qFLw9Bq8TkOMg8M8CcI92BYvFekXkjyCC9lXSEVa8Ck,2919
-brainstate/random.py,sha256=OMa4739GbrQpKspEx0TbeBqjA4yvwwJzdtpr-kJZN6s,187841
+brainstate/mixin.py,sha256=2f2toMUmgJIiovX1wi8OIBM8sbWH6s9Usa1ixL9J4tg,10747
+brainstate/mixin_test.py,sha256=-Ej9oUOu8O1M4oy37SVMj7xNRYhHHyAHwrjS_aISayo,2923
+brainstate/random.py,sha256=UbXfC0nrxk5FOsld0rCxBp2bIaeQHH5bj-NWQTR8bbQ,188447
 brainstate/random_test.py,sha256=cCeuYvlZkCS2_RgG0vipZFNSHG8b-uJ7SXM9SZDCYQM,17866
 brainstate/surrogate.py,sha256=1kgbn82GSlpReIytIVl29yozk75gkdZv0gTBlixQ4C4,43798
 brainstate/typing.py,sha256=tZyXpopNtp1UvqbntnmlUP6OBTM_ywyV9QBQ6Gu3IdU,2232
@@ -20,14 +20,14 @@ brainstate/functional/_others.py,sha256=ifB-l82y7ZB632yLUJOEcpkRY-yOoiJ0mtDOxNil
 brainstate/functional/_spikes.py,sha256=70qGvo4B--QtxfJMjLwGmk9pVsf2x2YNEEgjT-il_Jw,2574
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=jRTmfoUsH_315vW9YMZzyIn2KDAAsv56SplBnvOyBW0,1148
-brainstate/init/_generic.py,sha256=yLkmYRWpGHUK_nYx_fWPe74s_Cd-JDdQ6mUqpi4yqcc,7308
+brainstate/init/_generic.py,sha256=G5WHSQs3M9htREIZ_OXqj1ffSF_BdYdPTaMqe9AKj-k,7315
 brainstate/init/_random_inits.py,sha256=ycxH9WyKvPhRbk9PkamCLlc9aX5YokREBfSCbrbFQQ4,16021
 brainstate/init/_regular_inits.py,sha256=u77aSM0BkK9VULFJQZ1lIEYA_sJJzEZBTEttBSJ79RI,3090
 brainstate/nn/__init__.py,sha256=YJHoI8cXKVRS8f2vUl3Zegp5wm0svMz3qo9JmQJiMQk,2162
 brainstate/nn/_base.py,sha256=lzbZpku3Q2arH6ZaAwjs6bhbV0RcFChxo2UcpnX5t84,8481
 brainstate/nn/_connections.py,sha256=GSOW2IbpJRHdPyF4nFJ2RPgO8y6SVHT1Gn-pbri9pMk,22970
 brainstate/nn/_dynamics.py,sha256=OeYYXv1dqjUDcCsRhZo1XS7SP2li1vlH9uhME_PE9v0,13205
-brainstate/nn/_elementwise.py,sha256=6BTqSvSnaHhldwB5ol5OV0hPJ5yJ-Jpm4WSrtFKMNoQ,43579
+brainstate/nn/_elementwise.py,sha256=Br2yd1kdr06iWGSvpoebWWO6suXFDiF8PQv_hOX9kZQ,43599
 brainstate/nn/_embedding.py,sha256=WbgrIaM_14abN8zBDr0xipBOsFc8dXP2m7Z_aRLAfmU,2249
 brainstate/nn/_misc.py,sha256=Xc4U4NLmvfnKdBNDayFrRBPAy3p0beS6T9C59rIDP00,3790
 brainstate/nn/_normalizations.py,sha256=9yVDORAEpqEkL9MYSPU4m7C4q8Qj5UNsPh9sKmIt5gQ,14329
@@ -50,7 +50,7 @@ brainstate/optim/_sgd_optimizer.py,sha256=JiK_AVGregL0wn8uHhRQvK9Qq7Qja7dEyLW6Aa
 brainstate/transform/__init__.py,sha256=my2X4ZW0uKZRfN82zyGEPizWNJ0fsSP2akvmkjn43ck,1458
 brainstate/transform/_autograd.py,sha256=Pj_YxpU52guaxQs1NcB6qDtXgkvaPcoJbuvIF8T-Wmk,23964
 brainstate/transform/_autograd_test.py,sha256=RWriMemIF9FVFUjQh4IHzLhT9LGyd1JXpjXfFZKHn10,38654
-brainstate/transform/_control.py,sha256=NWceTIuLlj2uGTdNcqBAXgnaLuChOGgAtIXtFn5vdLU,26837
+brainstate/transform/_control.py,sha256=0NFUGLIenqKuBhBiTmY0YgCrl2GI1ZbuWMW0DSOolpE,26874
 brainstate/transform/_controls_test.py,sha256=mPUa_qmXXVxDziAJrPWRBwsGnc3cHR9co08eJB_fJwA,7648
 brainstate/transform/_jit.py,sha256=sjQHFV8Tt75fpdl12jjPRDPT92_IZxBBJAG4gapdbNQ,11471
 brainstate/transform/_jit_error.py,sha256=8rGRx8dtvmPWmHVOsfz30EUMXSix-m2PKM3Ni_9-_7I,4829
@@ -59,8 +59,8 @@ brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJ
 brainstate/transform/_make_jaxpr.py,sha256=vMPruKfp5Ugv8RL-9wGfQdSumLZdLtThZvv3sU9MDjE,30426
 brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
-brainstate-0.0.1.1.post20240708.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.1.1.post20240708.dist-info/METADATA,sha256=HN33Hoom47puNsauzxdmXvsmcP5RqrUpeR_2tVf7Y5U,3816
-brainstate-0.0.1.1.post20240708.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.1.1.post20240708.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.1.1.post20240708.dist-info/RECORD,,
+brainstate-0.0.1.1.post20240802.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.1.1.post20240802.dist-info/METADATA,sha256=LYHl7AJ94js5MhJVghX5mJ1KB_VA4kWVc1bCmZ0O3GY,3807
+brainstate-0.0.1.1.post20240802.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.0.1.1.post20240802.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.1.1.post20240802.dist-info/RECORD,,