brainstate 0.0.2.post20240913__py2.py3-none-any.whl → 0.0.2.post20241010__py2.py3-none-any.whl

brainstate/__init__.py

@@ -30,6 +30,8 @@ from . import surrogate
 from . import transform
 from . import typing
 from . import util
+from ._visualization import *
+from ._visualization import __all__ as _visualization_all
 from ._module import *
 from ._module import __all__ as _module_all
 from ._state import *
@@ -39,6 +41,6 @@ __all__ = (
     ['environ', 'share', 'nn', 'optim', 'random',
      'surrogate', 'functional', 'init',
      'mixin', 'transform', 'util', 'typing'] +
-    _module_all + _state_all
+    _module_all + _state_all + _visualization_all
 )
-del _module_all, _state_all
+del _module_all, _state_all, _visualization_all

brainstate/_module.py

@@ -56,14 +56,13 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
-from . import environ
-from ._state import State, StateDictManager, visible_state_dict
-from ._utils import set_module_as
-from .mixin import Mixin, Mode, DelayedInit, JointTypes, Batching, UpdateReturn
-from .transform import jit_error_if
-from .typing import Size, ArrayLike, PyTree
-from .util import unique_name, DictManager, get_unique_name
-
+from brainstate import environ
+from brainstate._state import State, StateDictManager, visible_state_dict
+from brainstate._utils import set_module_as
+from brainstate.mixin import Mixin, Mode, DelayedInit, JointTypes, Batching, UpdateReturn
+from brainstate.transform import jit_error_if
+from brainstate.typing import Size, ArrayLike, PyTree
+from brainstate.util import unique_name, DictManager, get_unique_name
 
 delay_identifier = '_*_delay_of_'
 _DELAY_ROTATE = 'rotation'
@@ -94,7 +93,7 @@ __all__ = [
 ]
 
 
-class Module(object):
+class Object:
   """
   The Module class for the whole ecosystem.
 
@@ -118,43 +117,9 @@ class Module(object):
   # the excluded nodes
   _invisible_nodes: Tuple[str, ...] = ()
 
-  # # the supported computing modes
-  # supported_modes: Optional[Sequence[Mode]] = None
-
-  def __init__(self, name: str = None, mode: Mode = None):
-    super().__init__()
-
-    # check whether the object has a unique name.
-    self._name = unique_name(self=self, name=name)
-
-    # mode setting
-    self._mode = None
-    self.mode = mode if mode is not None else environ.get('mode')
-
   def __repr__(self):
     return f'{self.__class__.__name__}'
 
-  @property
-  def name(self):
-    """Name of the model."""
-    return self._name
-
-  @name.setter
-  def name(self, name: str = None):
-    raise AttributeError('The name of the model is read-only.')
-
-  @property
-  def mode(self):
-    """Mode of the model, which is useful to control the multiple behaviors of the model."""
-    return self._mode
-
-  @mode.setter
-  def mode(self, value):
-    if not isinstance(value, Mode):
-      raise ValueError(f'Must be instance of {Mode.__name__}, '
-                       f'but we got {type(value)}: {value}')
-    self._mode = value
-
   def states(
       self,
       method: str = 'absolute',
@@ -238,30 +203,6 @@ class Module(object):
       nodes = nodes.unique()
     return nodes
 
-  def update(self, *args, **kwargs):
-    """
-    The function to specify the updating rule.
-    """
-    raise NotImplementedError(f'Subclass of {self.__class__.__name__} must '
-                              f'implement "update" function.')
-
-  def __call__(self, *args, **kwargs):
-    return self.update(*args, **kwargs)
-
-  def __rrshift__(self, other):
-    """
-    Support using right shift operator to call modules.
-
-    Examples
-    --------
-
-    >>> import brainstate as bst
-    >>> x = bst.random.rand((10, 10))
-    >>> l = bst.nn.Activation(jax.numpy.tanh)
-    >>> y = x >> l
-    """
-    return self.__call__(other)
-
   def init_state(self, *args, **kwargs):
     """
     State initialization function.
@@ -289,6 +230,23 @@ class Module(object):
     missing_keys = list(keys2 - keys1)
     return unexpected_keys, missing_keys
 
+  def __treescope_repr__(self, path, subtree_renderer):
+    import treescope  # type: ignore[import-not-found,import-untyped]
+    children = {}
+    for name, value in vars(self).items():
+      if name.startswith('_'):
+        continue
+      children[name] = value
+    return treescope.repr_lib.render_object_constructor(
+      object_type=type(self),
+      attributes=children,
+      path=path,
+      subtree_renderer=subtree_renderer,
+      color=treescope.formatting_util.color_from_string(
+        type(self).__qualname__
+      )
+    )
+
 
 def _find_nodes(self, method: str = 'absolute', level=-1, include_self=True, _lid=0, _edges=None) -> DictManager:
   if _edges is None:
@@ -374,6 +332,83 @@ def _add_node_relative(self, k, v, _paths, gather, nodes):
     nodes.append((k, v))
 
 
+class Module(Object):
+  """
+  The Module class for the whole ecosystem.
+
+  The ``Module`` is the base class for all the objects in the ecosystem. It
+  provides the basic functionalities for the objects, including:
+
+  - ``states()``: Collect all states in this node and the children nodes.
+  - ``nodes()``: Collect all children nodes.
+  - ``update()``: The function to specify the updating rule.
+  - ``init_state()``: State initialization function.
+  - ``save_state()``: Save states as a dictionary.
+  - ``load_state()``: Load states from the external objects.
+
+  """
+
+  __module__ = 'brainstate'
+
+  def __init__(self, name: str = None, mode: Mode = None):
+    super().__init__()
+
+    # check whether the object has a unique name.
+    self._name = unique_name(self=self, name=name)
+
+    # mode setting
+    self._mode = None
+    self.mode = mode if mode is not None else environ.get('mode')
+
+  def __repr__(self):
+    return f'{self.__class__.__name__}'
+
+  @property
+  def name(self):
+    """Name of the model."""
+    return self._name
+
+  @name.setter
+  def name(self, name: str = None):
+    raise AttributeError('The name of the model is read-only.')
+
+  @property
+  def mode(self):
+    """Mode of the model, which is useful to control the multiple behaviors of the model."""
+    return self._mode
+
+  @mode.setter
+  def mode(self, value):
+    if not isinstance(value, Mode):
+      raise ValueError(f'Must be instance of {Mode.__name__}, '
+                       f'but we got {type(value)}: {value}')
+    self._mode = value
+
+  def update(self, *args, **kwargs):
+    """
+    The function to specify the updating rule.
+    """
+    raise NotImplementedError(f'Subclass of {self.__class__.__name__} must '
+                              f'implement "update" function.')
+
+  def __call__(self, *args, **kwargs):
+    return self.update(*args, **kwargs)
+
+  def __rrshift__(self, other):
+    """
+    Support using right shift operator to call modules.
+
+    Examples
+    --------
+
+    >>> import brainstate as bst
+    >>> x = bst.random.rand((10, 10))
+    >>> l = bst.nn.Dropout(0.5)
+    >>> y = x >> l
+    """
+    return self.__call__(other)
+
+
 class Projection(Module):
   """
   Base class to model synaptic projections.

brainstate/_state.py

@@ -22,8 +22,8 @@ import numpy as np
 from jax.api_util import shaped_abstractify
 from jax.extend import source_info_util
 
-from .typing import ArrayLike, PyTree
-from .util import DictManager
+from brainstate.typing import ArrayLike, PyTree
+from brainstate.util import DictManager
 
 __all__ = [
   'State', 'ShortTermState', 'LongTermState', 'ParamState',

brainstate/_visualization.py

@@ -0,0 +1,47 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+
+__all__ = [
+  'display',
+]
+
+
+import importlib.util
+
+treescope_installed = importlib.util.find_spec('treescope') is not None
+try:
+  from IPython import get_ipython
+
+  in_ipython = get_ipython() is not None
+except ImportError:
+  in_ipython = False
+
+
+def display(*args):
+  """Display the given objects using the Treescope pretty-printer.
+
+  If treescope is not installed or the code is not running in IPython,
+  ``display`` will print the objects instead.
+  """
+  if not treescope_installed or not in_ipython:
+    for x in args:
+      print(x)
+    return
+
+  import treescope  # type: ignore[import-not-found,import-untyped]
+
+  for x in args:
+    treescope.display(x, ignore_exceptions=True, autovisualize=True)

brainstate/environ.py

@@ -16,12 +16,14 @@ from .mixin import Mode
 from .util import MemScaling, IdMemScaling
 
 __all__ = [
-  'set', 'context', 'get', 'all',
-  'set_host_device_count', 'set_platform',
-  'get_host_device_count', 'get_platform',
-  'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
-  'tolerance', 'register_default_behavior',
+  # functions for environment settings
+  'set', 'context', 'get', 'all', 'set_host_device_count', 'set_platform',
+  # functions for getting default behaviors
+  'get_host_device_count', 'get_platform', 'get_dt', 'get_mode', 'get_mem_scaling', 'get_precision',
+  # functions for default data types
   'dftype', 'ditype', 'dutype', 'dctype',
+  # others
+  'tolerance', 'register_default_behavior',
 ]
 
 # Default, there are several shared arguments in the global context.
@@ -57,8 +59,9 @@ def context(**kwargs):
 
   """
   if 'platform' in kwargs:
-    raise ValueError('Cannot set platform in environment context. '
-                     'Please use set_platform() or set() for the global setting.')
+    raise ValueError('\n'
+                     'Cannot set platform in "context" environment. \n'
+                     'You should set platform in the global environment by "set_platform()" or "set()".')
   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.')
@@ -135,6 +138,11 @@ def get(key: str, default: Any = _NOT_PROVIDE, desc: str = None):
 def all() -> dict:
   """
   Get all the current default computation environment.
+
+  Returns
+  -------
+  r: dict
+    The current default computation environment.
   """
   r = dict()
   for k, v in _environment_contexts.items():
@@ -227,6 +235,8 @@ def set(
   """
   Set the global default computation environment.
 
+
+
   Args:
     platform: str. The computing platform. Either 'cpu', 'gpu' or 'tpu'.
     host_device_count: int. The number of host devices.
@@ -290,6 +300,12 @@ def set_platform(platform: str):
   """
   Changes platform to CPU, GPU, or TPU. This utility only takes
   effect at the beginning of your program.
+
+  Args:
+    platform: str. The computing platform. Either 'cpu', 'gpu' or 'tpu'.
+
+  Raises:
+    ValueError: If the platform is not in ['cpu', 'gpu', 'tpu'].
   """
   assert platform in ['cpu', 'gpu', 'tpu']
   config.update("jax_platform_name", platform)
@@ -369,6 +385,23 @@ def _get_complex(precision: int):
 def dftype() -> DTypeLike:
   """
   Default floating data type.
+
+  This function returns the default floating data type based on the current precision.
+  If you want the data dtype is changed with the setting of the precision by ``brainstate.environ.set(precision)``,
+  you can use this function to get the default floating data type, and create the data by using ``dtype=dftype()``.
+
+  For example, if the precision is set to 32, the default floating data type is ``np.float32``.
+
+  >>> import brainstate as bst
+  >>> import numpy as np
+  >>> with bst.environ.context(precision=32):
+  ...    a = np.zeros(1, dtype=bst.environ.dftype())
+  >>> print(a.dtype)
+
+  Returns
+  -------
+  float_dtype: DTypeLike
+    The default floating data type.
   """
   return _get_float(get_precision())
 
@@ -376,6 +409,24 @@ def dftype() -> DTypeLike:
 def ditype() -> DTypeLike:
   """
   Default integer data type.
+
+  This function returns the default integer data type based on the current precision.
+  If you want the data dtype is changed with the setting of the precision by ``brainstate.environ.set(precision)``,
+  you can use this function to get the default integer data type, and create the data by using ``dtype=ditype()``.
+
+  For example, if the precision is set to 32, the default integer data type is ``np.int32``.
+
+  >>> import brainstate as bst
+  >>> import numpy as np
+  >>> with bst.environ.context(precision=32):
+  ...    a = np.zeros(1, dtype=bst.environ.ditype())
+  >>> print(a.dtype)
+  int32
+
+  Returns
+  -------
+  int_dtype: DTypeLike
+    The default integer data type.
   """
   return _get_int(get_precision())
 
@@ -383,6 +434,25 @@ def ditype() -> DTypeLike:
 def dutype() -> DTypeLike:
   """
   Default unsigned integer data type.
+
+  This function returns the default unsigned integer data type based on the current precision.
+  If you want the data dtype is changed with the setting of the precision
+  by ``brainstate.environ.set(precision)``, you can use this function to get the default
+  unsigned integer data type, and create the data by using ``dtype=dutype()``.
+
+  For example, if the precision is set to 32, the default unsigned integer data type is ``np.uint32``.
+
+  >>> import brainstate as bst
+  >>> import numpy as np
+  >>> with bst.environ.context(precision=32):
+  ...   a = np.zeros(1, dtype=bst.environ.dutype())
+  >>> print(a.dtype)
+  uint32
+
+  Returns
+  -------
+  uint_dtype: DTypeLike
+    The default unsigned integer data type.
   """
   return _get_uint(get_precision())
 
@@ -390,6 +460,24 @@ def dutype() -> DTypeLike:
 def dctype() -> DTypeLike:
   """
   Default complex data type.
+
+  This function returns the default complex data type based on the current precision.
+  If you want the data dtype is changed with the setting of the precision by ``brainstate.environ.set(precision)``,
+  you can use this function to get the default complex data type, and create the data by using ``dtype=dctype()``.
+
+  For example, if the precision is set to 32, the default complex data type is ``np.complex64``.
+
+  >>> import brainstate as bst
+  >>> import numpy as np
+  >>> with bst.environ.context(precision=32):
+  ...    a = np.zeros(1, dtype=bst.environ.dctype())
+  >>> print(a.dtype)
+  complex64
+
+  Returns
+  -------
+  complex_dtype: DTypeLike
+    The default complex data type.
   """
   return _get_complex(get_precision())
 
@@ -405,7 +493,27 @@ def tolerance():
 
 def register_default_behavior(key: str, behavior: Callable, replace_if_exist: bool = False):
   """
-  Register a default behavior for a specific key.
+  Register a default behavior for a specific global key parameter.
+
+  For example, you can register a default behavior for the key 'dt' by::
+
+  >>> import brainstate as bst
+  >>> def dt_behavior(dt):
+  ...     print(f'Set the default dt to {dt}.')
+  ...
+  >>> bst.environ.register_default_behavior('dt', dt_behavior)
+
+  Then, when you set the default dt by `brainstate.environ.set(dt=0.1)`, the behavior
+  `dt_behavior` will be called with
+  `dt_behavior(0.1)`.
+
+  >>> bst.environ.set(dt=0.1)
+  Set the default dt to 0.1.
+  >>> with bst.environ.context(dt=0.2):
+  ...     pass
+  Set the default dt to 0.2.
+  Set the default dt to 0.1.
+
 
   Args:
     key: str. The key to register.
@@ -421,4 +529,3 @@ def register_default_behavior(key: str, behavior: Callable, replace_if_exist: bo
 
 
 set(dt=0.1, precision=32, mode=Mode(), mem_scaling=IdMemScaling())
-

brainstate/environ_test.py

@@ -0,0 +1,56 @@
+# Copyright 2024 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.
+# ==============================================================================
+
+
+import unittest
+
+import jax.numpy as jnp
+
+import brainstate as bst
+
+
+class TestEnviron(unittest.TestCase):
+  def test_precision(self):
+    with bst.environ.context(precision=64):
+      a = bst.random.randn(1)
+      self.assertEqual(a.dtype, jnp.float64)
+
+    with bst.environ.context(precision=32):
+      a = bst.random.randn(1)
+      self.assertEqual(a.dtype, jnp.float32)
+
+    with bst.environ.context(precision=16):
+      a = bst.random.randn(1)
+      self.assertEqual(a.dtype, jnp.bfloat16)
+
+  def test_platform(self):
+    with self.assertRaises(ValueError):
+      with bst.environ.context(platform='cpu'):
+        a = bst.random.randn(1)
+        self.assertEqual(a.device(), 'cpu')
+
+  def test_register_default_behavior(self):
+    dt_ = 0.1
+
+    def dt_behavior(dt):
+      nonlocal dt_
+      dt_ = dt
+      print(f'dt: {dt}')
+
+    bst.environ.register_default_behavior('dt', dt_behavior)
+
+    with bst.environ.context(dt=0.2):
+      self.assertEqual(dt_, 0.2)
+    self.assertEqual(dt_, 0.1)