brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/_module_test.py

@@ -1,207 +0,0 @@
-# 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 jaxlib.xla_extension
-
-import brainstate as bst
-
-
-class TestDelay(unittest.TestCase):
-  def test_delay1(self):
-    a = bst.State(bst.random.random(10, 20))
-    delay = bst.Delay(a.value)
-    delay.register_entry('a', 1.)
-    delay.register_entry('b', 2.)
-    delay.register_entry('c', None)
-
-    delay.init_state()
-    with self.assertRaises(KeyError):
-      delay.register_entry('c', 10.)
-    bst.util.clear_buffer_memory()
-
-  def test_rotation_delay(self):
-    rotation_delay = bst.Delay(jnp.ones((1,)))
-    t0 = 0.
-    t1, n1 = 1., 10
-    t2, n2 = 2., 20
-
-    rotation_delay.register_entry('a', t0)
-    rotation_delay.register_entry('b', t1)
-    rotation_delay.register_entry('c2', 1.9)
-    rotation_delay.register_entry('c', t2)
-
-    rotation_delay.init_state()
-
-    print()
-    # print(rotation_delay)
-    # print(rotation_delay.max_length)
-
-    for i in range(100):
-      bst.environ.set(i=i)
-      rotation_delay(jnp.ones((1,)) * i)
-      # print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c2'), rotation_delay.at('c'))
-      self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
-      self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
-      self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-    bst.util.clear_buffer_memory()
-
-  def test_concat_delay(self):
-    rotation_delay = bst.Delay(jnp.ones([1]), delay_method='concat')
-    t0 = 0.
-    t1, n1 = 1., 10
-    t2, n2 = 2., 20
-
-    rotation_delay.register_entry('a', t0)
-    rotation_delay.register_entry('b', t1)
-    rotation_delay.register_entry('c', t2)
-
-    rotation_delay.init_state()
-
-    print()
-    for i in range(100):
-      bst.environ.set(i=i)
-      rotation_delay(jnp.ones((1,)) * i)
-      print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c'))
-      self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
-      self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
-      self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-    bst.util.clear_buffer_memory()
-
-  def test_jit_erro(self):
-    rotation_delay = bst.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
-    rotation_delay.init_state()
-
-    with bst.environ.context(i=0, t=0, jit_error_check=True):
-      rotation_delay.retrieve_at_time(-2.0)
-      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-        rotation_delay.retrieve_at_time(-2.1)
-      rotation_delay.retrieve_at_time(-2.01)
-      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-        rotation_delay.retrieve_at_time(-2.09)
-      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-        rotation_delay.retrieve_at_time(0.1)
-      with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-        rotation_delay.retrieve_at_time(0.01)
-
-  def test_round_interp(self):
-    for shape in [(1,), (1, 1), (1, 1, 1)]:
-      for delay_method in ['rotation', 'concat']:
-        rotation_delay = bst.Delay(jnp.ones(shape), time=2., delay_method=delay_method, interp_method='round')
-        t0, n1 = 0.01, 0
-        t1, n1 = 1.04, 10
-        t2, n2 = 1.06, 11
-        rotation_delay.init_state()
-
-        @bst.transform.jit
-        def retrieve(td, i):
-          with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
-            return rotation_delay.retrieve_at_time(td)
-
-        print()
-        for i in range(100):
-          t = i * bst.environ.get_dt()
-          with bst.environ.context(i=i, t=t):
-            rotation_delay(jnp.ones(shape) * i)
-            print(i,
-                  retrieve(t - t0, i),
-                  retrieve(t - t1, i),
-                  retrieve(t - t2, i))
-            self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.ones(shape) * i))
-            self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
-            self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
-        bst.util.clear_buffer_memory()
-
-  def test_linear_interp(self):
-    for shape in [(1,), (1, 1), (1, 1, 1)]:
-      for delay_method in ['rotation', 'concat']:
-        print(shape, delay_method)
-
-        rotation_delay = bst.Delay(jnp.ones(shape), time=2., delay_method=delay_method, interp_method='linear_interp')
-        t0, n0 = 0.01, 0.1
-        t1, n1 = 1.04, 10.4
-        t2, n2 = 1.06, 10.6
-        rotation_delay.init_state()
-
-        @bst.transform.jit
-        def retrieve(td, i):
-          with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
-            return rotation_delay.retrieve_at_time(td)
-
-        print()
-        for i in range(100):
-          t = i * bst.environ.get_dt()
-          with bst.environ.context(i=i, t=t):
-            rotation_delay(jnp.ones(shape) * i)
-            print(i,
-                  retrieve(t - t0, i),
-                  retrieve(t - t1, i),
-                  retrieve(t - t2, i))
-            self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.maximum(jnp.ones(shape) * i - n0, 0.)))
-            self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
-            self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
-        bst.util.clear_buffer_memory()
-
-  def test_rotation_and_concat_delay(self):
-    rotation_delay = bst.Delay(jnp.ones((1,)))
-    concat_delay = bst.Delay(jnp.ones([1]), delay_method='concat')
-    t0 = 0.
-    t1, n1 = 1., 10
-    t2, n2 = 2., 20
-
-    rotation_delay.register_entry('a', t0)
-    rotation_delay.register_entry('b', t1)
-    rotation_delay.register_entry('c', t2)
-    concat_delay.register_entry('a', t0)
-    concat_delay.register_entry('b', t1)
-    concat_delay.register_entry('c', t2)
-
-    rotation_delay.init_state()
-    concat_delay.init_state()
-
-    print()
-    for i in range(100):
-      bst.environ.set(i=i)
-      new = jnp.ones((1,)) * i
-      rotation_delay(new)
-      concat_delay(new)
-      self.assertTrue(jnp.allclose(rotation_delay.at('a'), concat_delay.at('a'), ))
-      self.assertTrue(jnp.allclose(rotation_delay.at('b'), concat_delay.at('b'), ))
-      self.assertTrue(jnp.allclose(rotation_delay.at('c'), concat_delay.at('c'), ))
-    bst.util.clear_buffer_memory()
-
-
-class TestModule(unittest.TestCase):
-  def test_states(self):
-    class A(bst.Module):
-      def __init__(self):
-        super().__init__()
-        self.a = bst.State(bst.random.random(10, 20))
-        self.b = bst.State(bst.random.random(10, 20))
-
-    class B(bst.Module):
-      def __init__(self):
-        super().__init__()
-        self.a = A()
-        self.b = bst.State(bst.random.random(10, 20))
-
-    b = B()
-    print()
-    print(b.states())
-    print(b.states())
-    print(b.states(level=0))
-    print(b.states(level=0))

brainstate/nn/_base.py

@@ -1,251 +0,0 @@
-# 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.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-from __future__ import annotations
-
-import inspect
-from typing import Sequence, Optional, Tuple, Union
-
-from brainstate._module import Module, UpdateReturn, Container, visible_module_dict
-from brainstate.mixin import Mixin, DelayedInitializer, DelayedInit
-
-__all__ = [
-  'ExplicitInOutSize',
-  'ElementWiseBlock',
-  'Sequential',
-  'DnnLayer',
-]
-
-
-# -------------------------------------------------------------------------------------- #
-# Network Related Concepts
-# -------------------------------------------------------------------------------------- #
-
-
-class ExplicitInOutSize(Mixin):
-  """
-  Mix-in class with the explicit input and output shape.
-
-  Attributes
-  ----------
-  in_size: tuple[int]
-    The input shape, without the batch size. This argument is important, since it is
-    used to evaluate the shape of the output.
-  out_size: tuple[int]
-    The output shape, without the batch size.
-  """
-  __module__ = 'brainstate.nn'
-
-  _in_size: Optional[Tuple[int, ...]] = None
-  _out_size: Optional[Tuple[int, ...]] = None
-
-  @property
-  def in_size(self) -> Tuple[int, ...]:
-    return self._in_size
-
-  @in_size.setter
-  def in_size(self, in_size: Sequence[int] | int):
-    if isinstance(in_size, int):
-      in_size = (in_size,)
-    assert isinstance(in_size, (tuple, list)), f"Invalid type of in_size: {type(in_size)}"
-    self._in_size = tuple(in_size)
-
-  @property
-  def out_size(self) -> Tuple[int, ...]:
-    return self._out_size
-
-  @out_size.setter
-  def out_size(self, out_size: Sequence[int] | int):
-    if isinstance(out_size, int):
-      out_size = (out_size,)
-    assert isinstance(out_size, (tuple, list)), f"Invalid type of out_size: {type(out_size)}"
-    self._out_size = tuple(out_size)
-
-
-class ElementWiseBlock(Mixin):
-  """
-  Mix-in class for element-wise modules.
-  """
-  __module__ = 'brainstate.nn'
-
-
-class Sequential(Module, UpdateReturn, Container, ExplicitInOutSize):
-  """
-  A sequential `input-output` module.
-
-  Modules will be added to it in the order they are passed in the
-  constructor. Alternatively, an ``dict`` of modules can be
-  passed in. The ``update()`` method of ``Sequential`` accepts any
-  input and forwards it to the first module it contains. It then
-  "chains" outputs to inputs sequentially for each subsequent module,
-  finally returning the output of the last module.
-
-  The value a ``Sequential`` provides over manually calling a sequence
-  of modules is that it allows treating the whole container as a
-  single module, such that performing a transformation on the
-  ``Sequential`` applies to each of the modules it stores (which are
-  each a registered submodule of the ``Sequential``).
-
-  What's the difference between a ``Sequential`` and a
-  :py:class:`Container`? A ``Container`` is exactly what it
-  sounds like--a container to store :py:class:`DynamicalSystem` s!
-  On the other hand, the layers in a ``Sequential`` are connected
-  in a cascading way.
-
-  Examples
-  --------
-
-  >>> import jax
-  >>> import brainstate as bst
-  >>> import brainstate.nn as nn
-  >>>
-  >>> # composing ANN models
-  >>> l = nn.Sequential(nn.Linear(100, 10),
-  >>>                   jax.nn.relu,
-  >>>                   nn.Linear(10, 2))
-  >>> l(bst.random.random((256, 100)))
-  >>>
-  >>> # Using Sequential with Dict. This is functionally the
-  >>> # same as the above code
-  >>> l = nn.Sequential(l1=nn.Linear(100, 10),
-  >>>                   l2=jax.nn.relu,
-  >>>                   l3=nn.Linear(10, 2))
-  >>> l(bst.random.random((256, 100)))
-
-  Args:
-    modules_as_tuple: The children modules.
-    modules_as_dict: The children modules.
-    name: The object name.
-    mode: The object computing context/mode. Default is ``None``.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self, first: ExplicitInOutSize, *modules_as_tuple, **modules_as_dict):
-    super().__init__()
-
-    assert isinstance(first, ExplicitInOutSize)
-    in_size = first.out_size
-
-    tuple_modules = []
-    for module in modules_as_tuple:
-      module, in_size = self._format_module(module, in_size)
-      tuple_modules.append(module)
-
-    dict_modules = dict()
-    for key, module in modules_as_dict.items():
-      module, in_size = self._format_module(module, in_size)
-      dict_modules[key] = module
-
-    # Attribute of "Container"
-    self.children = visible_module_dict(self.format_elements(object, first, *tuple_modules, **dict_modules))
-
-    # the input and output shape
-    if first.in_size is not None:
-      self.in_size = first.in_size
-    self.out_size = tuple(in_size)
-
-  def _format_module(self, module, in_size):
-    if isinstance(module, DelayedInitializer):
-      module = module(in_size=in_size)
-      assert isinstance(module, ExplicitInOutSize)
-      out_size = module.out_size
-    elif isinstance(module, ElementWiseBlock):
-      out_size = in_size
-    elif isinstance(module, ExplicitInOutSize):
-      out_size = module.out_size
-    else:
-      raise TypeError(f"Unsupported type {type(module)}. ")
-    return module, out_size
-
-  def update(self, x):
-    """Update function of a sequential model.
-    """
-    for m in self.children.values():
-      x = m(x)
-    return x
-
-  def update_return(self):
-    """
-    The return information of the sequence according to the final model.
-    """
-    last = self[-1]
-    if not isinstance(last, UpdateReturn):
-      raise NotImplementedError(f'The last element in the sequence is not an instance of {UpdateReturn.__name__}')
-    return last.update_return()
-
-  def update_return_info(self):
-    """
-    The return information of the sequence according to the final model.
-    """
-    last = self[-1]
-    if not isinstance(last, UpdateReturn):
-      raise NotImplementedError(f'The last element in the sequence is not an instance of {UpdateReturn.__name__}')
-    return last.update_return_info()
-
-  def __getitem__(self, key: Union[int, slice, str]):
-    if isinstance(key, str):
-      if key in self.children:
-        return self.children[key]
-      else:
-        raise KeyError(f'Does not find a component named {key} in\n {str(self)}')
-    elif isinstance(key, slice):
-      return Sequential(**dict(tuple(self.children.items())[key]))
-    elif isinstance(key, int):
-      return tuple(self.children.values())[key]
-    elif isinstance(key, (tuple, list)):
-      _all_nodes = tuple(self.children.items())
-      return Sequential(**dict(_all_nodes[k] for k in key))
-    else:
-      raise KeyError(f'Unknown type of key: {type(key)}')
-
-  def __repr__(self):
-    nodes = self.children.values()
-    entries = '\n'.join(f'  [{i}] {_repr_object(x)}' for i, x in enumerate(nodes))
-    return f'{self.__class__.__name__}(\n{entries}\n)'
-
-
-def _repr_object(x):
-  if isinstance(x, Module):
-    return repr(x)
-  elif callable(x):
-    signature = inspect.signature(x)
-    args = [f'{k}={v.default}' for k, v in signature.parameters.items()
-            if v.default is not inspect.Parameter.empty]
-    args = ', '.join(args)
-    while not hasattr(x, '__name__'):
-      if not hasattr(x, 'func'):
-        break
-      x = x.func  # Handle functools.partial
-    if not hasattr(x, '__name__') and hasattr(x, '__class__'):
-      return x.__class__.__name__
-    if args:
-      return f'{x.__name__}(*, {args})'
-    return x.__name__
-  else:
-    x = repr(x).split('\n')
-    x = [x[0]] + ['  ' + y for y in x[1:]]
-    return '\n'.join(x)
-
-
-class DnnLayer(Module, ExplicitInOutSize, DelayedInit):
-  """
-  A DNN layer.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __repr__(self):
-    return f"{self.__class__.__name__}(in_size={self.in_size}, out_size={self.out_size})"