Bump version to 2.7.4 and update environment handling in brainstate integration

Author: chaoming0625

brainpy/__init__.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "2.7.3"
-__version_info__ = (2, 7, 3)
+__version__ = "2.7.4"
+__version_info__ = tuple(map(int, __version__.split(".")))
 
 
 from brainpy import _errors as errors
@@ -142,7 +142,6 @@
     ArrayCollector as ArrayCollector,
     Collector as Collector,
 )
-from brainpy import state
 
 from brainpy.deprecations import deprecation_getattr
 
@@ -151,7 +150,6 @@
 
 
 if __name__ == '__main__':
-    state
     connect
     initialize,  # weight initialization
     optim,  # gradient descent optimizers

brainpy/context.py

@@ -21,6 +21,7 @@
 from typing import Any, Union
 
 import brainstate
+from brainpy.math.defaults import env
 from brainpy.tools.dicts import DotDict
 
 __all__ = [
@@ -40,14 +41,14 @@ def __init__(self):
 
     @property
     def dt(self):
-        return brainstate.environ.get_dt()
+        return brainstate.environ.get_dt(env=env)
 
     @dt.setter
     def dt(self, dt):
         self.set_dt(dt)
 
     def set_dt(self, dt: Union[int, float]):
-        brainstate.environ.set(dt=dt)
+        brainstate.environ.set(dt=dt, env=env)
 
     def load(self, key, value: Any = None, desc: str = None):
         """Load the shared data by the ``key``.
@@ -57,16 +58,16 @@ def load(self, key, value: Any = None, desc: str = None):
           value (Any): the default value when ``key`` is not defined in the shared.
           desc: (str): the description of the key.
         """
-        return brainstate.environ.get(key, value, desc)
+        return brainstate.environ.get(key, value, desc, env=env)
 
     def save(self, *args, **kwargs) -> None:
         """Save shared arguments in the global context."""
         assert len(args) % 2 == 0
         for i in range(0, len(args), 2):
             identifier = args[i]
             data = args[i + 1]
-            brainstate.environ.set(**{identifier: data})
-        brainstate.environ.set(**kwargs)
+            brainstate.environ.set(**{identifier: data}, env=env)
+        brainstate.environ.set(**kwargs, env=env)
 
     def __setitem__(self, key, value):
         """Enable setting the shared item by ``bp.share[key] = value``."""
@@ -78,7 +79,7 @@ def __getitem__(self, item):
 
     def get_shargs(self) -> DotDict:
         """Get all shared arguments in the global context."""
-        return DotDict(brainstate.environ.all())
+        return DotDict(brainstate.environ.all(env=env))
 
 
 share = _ShareContext()

brainpy/math/defaults.py

@@ -19,6 +19,8 @@
 from .modes import NonBatchingMode
 from .scales import IdScaling
 
+env = brainstate.environ.EnvironmentState()
+
 
 class setting:
     def __init__(self):
@@ -42,79 +44,80 @@ def __init__(self):
             # default return array type
             # numpy_func_return='jax_array',  # 'bp_array','jax_array'
             numpy_func_return='bp_array',  # 'bp_array','jax_array'
+            env=env,
         )
 
     @property
     def mode(self):
-        return brainstate.environ.get('mode')
+        return brainstate.environ.get('mode', env=env)
 
     @property
     def membrane_scaling(self):
-        return brainstate.environ.get('membrane_scaling')
+        return brainstate.environ.get('membrane_scaling', env=env)
 
     @property
     def dt(self):
-        return brainstate.environ.get('dt')
+        return brainstate.environ.get('dt', env=env)
 
     @property
     def bool_(self):
-        return brainstate.environ.get('bool_')
+        return brainstate.environ.get('bool_', env=env)
 
     @property
     def int_(self):
-        return brainstate.environ.get('int_')
+        return brainstate.environ.get('int_', env=env)
 
     @property
     def float_(self):
-        return brainstate.environ.get('float_')
+        return brainstate.environ.get('float_', env=env)
 
     @property
     def complex_(self):
-        return brainstate.environ.get('complex_')
+        return brainstate.environ.get('complex_', env=env)
 
     @property
     def bp_object_as_pytree(self):
-        return brainstate.environ.get('bp_object_as_pytree')
+        return brainstate.environ.get('bp_object_as_pytree', env=env)
 
     @property
     def numpy_func_return(self):
-        return brainstate.environ.get('numpy_func_return')
+        return brainstate.environ.get('numpy_func_return', env=env)
 
     @mode.setter
     def mode(self, value):
-        brainstate.environ.set(mode=value)
+        brainstate.environ.set(mode=value, env=env)
 
     @membrane_scaling.setter
     def membrane_scaling(self, value):
-        brainstate.environ.set(membrane_scaling=value)
+        brainstate.environ.set(membrane_scaling=value, env=env)
 
     @dt.setter
     def dt(self, value):
-        brainstate.environ.set(dt=value)
+        brainstate.environ.set(dt=value, env=env)
 
     @bool_.setter
     def bool_(self, value):
-        brainstate.environ.set(bool_=value)
+        brainstate.environ.set(bool_=value, env=env)
 
     @int_.setter
     def int_(self, value):
-        brainstate.environ.set(int_=value)
+        brainstate.environ.set(int_=value, env=env)
 
     @float_.setter
     def float_(self, value):
-        brainstate.environ.set(float_=value)
+        brainstate.environ.set(float_=value, env=env)
 
     @complex_.setter
     def complex_(self, value):
-        brainstate.environ.set(complex_=value)
+        brainstate.environ.set(complex_=value, env=env)
 
     @bp_object_as_pytree.setter
     def bp_object_as_pytree(self, value):
-        brainstate.environ.set(bp_object_as_pytree=value)
+        brainstate.environ.set(bp_object_as_pytree=value, env=env)
 
     @numpy_func_return.setter
     def numpy_func_return(self, value):
-        brainstate.environ.set(numpy_func_return=value)
+        brainstate.environ.set(numpy_func_return=value, env=env)
 
 
 defaults = setting()

brainpy/state/__init__.py

@@ -43,7 +43,6 @@
 
 sys.path.insert(0, os.path.abspath('../docs/'))
 sys.path.insert(0, os.path.abspath('../'))
-os.makedirs('../docs/_static', exist_ok=True)
 shutil.copytree('../images/', './_static/logos/', dirs_exist_ok=True)
 shutil.copyfile('../changelog.md', './changelog.md')