brainstate 0.0.2.post20240814__py2.py3-none-any.whl → 0.0.2.post20240825__py2.py3-none-any.whl

brainstate/environ.py

@@ -6,7 +6,7 @@ import functools
 import os
 import re
 from collections import defaultdict
-from typing import Any
+from typing import Any, Callable
 
 import numpy as np
 from jax import config, devices, numpy as jnp
@@ -20,7 +20,7 @@ __all__ = [
   'set_host_device_count', 'set_platform',
   'get_host_device_count', 'get_platform',
   'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
-  'tolerance',
+  'tolerance', 'register_default_behavior',
   'dftype', 'ditype', 'dutype', 'dctype',
 ]
 
@@ -31,8 +31,9 @@ JIT_ERROR_CHECK = 'jit_error_check'  # whether to record the current computation
 FIT = 'fit'  # whether to fit the model.
 
 _NOT_PROVIDE = object()
-_environment_defaults = dict()
-_environment_contexts = defaultdict(list)
+_environment_defaults = dict()  # default environment settings
+_environment_contexts = defaultdict(list)  # current environment settings
+_environment_functions = dict()  # environment functions
 
 
 @contextlib.contextmanager
@@ -61,19 +62,34 @@ def context(**kwargs):
   if 'host_device_count' in kwargs:
     raise ValueError('Cannot set host_device_count in environment context. '
                      'Please use set_host_device_count() or set() for the global setting.')
+
   if 'precision' in kwargs:
     last_precision = get_precision()
     _set_jax_precision(kwargs['precision'])
 
   try:
-    # update the current environment
     for k, v in kwargs.items():
+
+      # update the current environment
       _environment_contexts[k].append(v)
+
+      # restore the environment functions
+      if k in _environment_functions:
+        _environment_functions[k](v)
+
     # yield the current all environment information
     yield all()
   finally:
+
     for k, v in kwargs.items():
+
+      # restore the current environment
       _environment_contexts[k].pop()
+
+      # restore the environment functions
+      if k in _environment_functions:
+        _environment_functions[k](get(k))
+
     if 'precision' in kwargs:
       _set_jax_precision(last_precision)
 
@@ -232,8 +248,15 @@ def set(
   if mode is not None:
     assert isinstance(mode, Mode), 'mode must be a Mode instance.'
     kwargs['mode'] = mode
+
+  # set default environment
   _environment_defaults.update(kwargs)
 
+  # update the environment functions
+  for k, v in kwargs.items():
+    if k in _environment_functions:
+      _environment_functions[k](v)
+
 
 def set_host_device_count(n):
   """
@@ -258,6 +281,10 @@ def set_host_device_count(n):
   xla_flags = re.sub(r"--xla_force_host_platform_device_count=\S+", "", xla_flags).split()
   os.environ["XLA_FLAGS"] = " ".join(["--xla_force_host_platform_device_count={}".format(n)] + xla_flags)
 
+  # update the environment functions
+  if 'host_device_count' in _environment_functions:
+    _environment_functions['host_device_count'](n)
+
 
 def set_platform(platform: str):
   """
@@ -267,6 +294,10 @@ def set_platform(platform: str):
   assert platform in ['cpu', 'gpu', 'tpu']
   config.update("jax_platform_name", platform)
 
+  # update the environment functions
+  if 'platform' in _environment_functions:
+    _environment_functions['platform'](platform)
+
 
 def _set_jax_precision(precision: int):
   """
@@ -372,4 +403,22 @@ def tolerance():
     return jnp.array(1e-2, dtype=np.float16)
 
 
+def register_default_behavior(key: str, behavior: Callable, replace_if_exist: bool = False):
+  """
+  Register a default behavior for a specific key.
+
+  Args:
+    key: str. The key to register.
+    behavior: Callable. The behavior to register. It should be a callable.
+    replace_if_exist: bool. Whether to replace the behavior if the key has been registered.
+
+  """
+  assert isinstance(key, str), 'key must be a string.'
+  assert callable(behavior), 'behavior must be a callable.'
+  if not replace_if_exist:
+    assert key not in _environment_functions, f'{key} has been registered.'
+  _environment_functions[key] = behavior
+
+
 set(dt=0.1, precision=32, mode=Mode(), mem_scaling=IdMemScaling())
+

{brainstate-0.0.2.post20240814.dist-info → brainstate-0.0.2.post20240825.dist-info}/METADATA

@@ -1,14 +1,14 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.0.2.post20240814
+Version: 0.0.2.post20240825
 Summary: A State-based Transformation System for Brain Dynamics Programming.
 Home-page: https://github.com/brainpy/brainstate
 Author: BDP
-Author-email: BrainPy Team <chao.brain@qq.com>
+Author-email: BrainState Developers <chao.brain@qq.com>
 License: Apache-2.0 license
 Project-URL: homepage, http://github.com/brainpy
 Project-URL: repository, http://github.com/brainpy/brainstate
-Keywords: brainpy,brain simulation,brain-inspired computing
+Keywords: computational neuroscience,brain-inspired computation,brain dynamics programming
 Classifier: Natural Language :: English
 Classifier: Operating System :: OS Independent
 Classifier: Development Status :: 4 - Beta
@@ -31,11 +31,11 @@ License-File: LICENSE
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
-Requires-Dist: brainunit
+Requires-Dist: brainunit >=0.0.2
 Provides-Extra: cpu
 Requires-Dist: jaxlib ; extra == 'cpu'
 Provides-Extra: cuda12
-Requires-Dist: jaxlib[cuda12_pip] ; extra == 'cuda12'
+Requires-Dist: jaxlib[cuda12] ; extra == 'cuda12'
 Provides-Extra: testing
 Requires-Dist: pytest ; extra == 'testing'
 Provides-Extra: tpu

{brainstate-0.0.2.post20240814.dist-info → brainstate-0.0.2.post20240825.dist-info}/RECORD

@@ -5,7 +5,7 @@ brainstate/_random_for_unit.py,sha256=1rHr7gfH_bYrJfpxbDhQUk_j00Yosx-GzyZCXrLxsd
 brainstate/_state.py,sha256=C0widCOj_ca6zfqh95jzFXf_G5vi0hJyuQ5GIqEqOUs,12102
 brainstate/_state_test.py,sha256=HDdipndRLhEHWEdTmyT1ayEBkbv6qJKykfCWKI6yJ_E,1253
 brainstate/_utils.py,sha256=RLorgGJkt2BhbX4C-ygd-PPG0wfcGCghjSP93sRvzqM,833
-brainstate/environ.py,sha256=LwRwnFaTbv8l7nHRIbSV46WzcN7pGLQFhT_xDUox2yA,10240
+brainstate/environ.py,sha256=k0p1oyi9jbsPfuvqrPL-_zgSd7VW3LRs0LboxlaaIfc,11806
 brainstate/mixin.py,sha256=OumTTSVyYSbtudjfS_MRThsBaeVJ_0JggeMClY7xtBA,10758
 brainstate/mixin_test.py,sha256=-Ej9oUOu8O1M4oy37SVMj7xNRYhHHyAHwrjS_aISayo,2923
 brainstate/random.py,sha256=BqEBYVD9TGe8dSzp8U0suK0O4r6Ox59GCq0mwfUndVQ,188073
@@ -59,8 +59,8 @@ brainstate/transform/_jit_test.py,sha256=5ltT7izh_OS9dcHnRymmVhq01QomjwZGdA8XzwJ
 brainstate/transform/_make_jaxpr.py,sha256=ZkrOZu4_0xcILuPUA3RFEkorJ-xbDuDtXorJI_qVThE,30450
 brainstate/transform/_make_jaxpr_test.py,sha256=K3vRUBroDTCCx0lnmhgHtgrlWvWglJO2f1K2phTvU70,3819
 brainstate/transform/_progress_bar.py,sha256=VGoRZPRBmB8ELNwLc6c7S8QhUUTvn0FY46IbBm9cuYM,3502
-brainstate-0.0.2.post20240814.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.0.2.post20240814.dist-info/METADATA,sha256=skMWlfxiGaJHxzQS7dY95V91umhXVjN_HuhGO0xHP1M,3805
-brainstate-0.0.2.post20240814.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.0.2.post20240814.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.0.2.post20240814.dist-info/RECORD,,
+brainstate-0.0.2.post20240825.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.0.2.post20240825.dist-info/METADATA,sha256=COfpxoCL7w1xGa1OFYFeANFLhAKmSioWVtmF_i2st34,3849
+brainstate-0.0.2.post20240825.dist-info/WHEEL,sha256=GUeE9LxUgRABPG7YM0jCNs9cBsAIx0YAkzCB88PMLgc,109
+brainstate-0.0.2.post20240825.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.0.2.post20240825.dist-info/RECORD,,

{brainstate-0.0.2.post20240814.dist-info → brainstate-0.0.2.post20240825.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.38.4)
+Generator: setuptools (73.0.1)
 Root-Is-Purelib: true
 Tag: py2-none-any
 Tag: py3-none-any