brainstate 0.1.2__py2.py3-none-any.whl → 0.1.3__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 from . import augment
 from . import compile

brainstate/_compatible_import.py

@@ -19,7 +19,7 @@
 import importlib.util
 from contextlib import contextmanager
 from functools import partial
-from typing import Iterable, Hashable, TypeVar, Callable
+from typing import Iterable, Hashable, TypeVar, Callable, TYPE_CHECKING
 
 import jax
 
@@ -45,8 +45,8 @@ T1 = TypeVar("T1")
 T2 = TypeVar("T2")
 T3 = TypeVar("T3")
 
-
 from saiunit._compatible_import import wrap_init
+
 brainevent_installed = importlib.util.find_spec('brainevent') is not None
 
 from jax.core import get_aval, Tracer
@@ -151,13 +151,13 @@ def to_concrete_aval(aval):
     return aval
 
 
-if brainevent_installed:
-    import brainevent
-else:
+if not brainevent_installed:
+    if not TYPE_CHECKING:
+        class BrainEvent:
+            def __getattr__(self, item):
+                raise ImportError('brainevent is not installed, please install brainevent first.')
 
-    class BrainEvent:
-        def __getattr__(self, item):
-            raise ImportError('brainevent is not installed, please install brainevent first.')
+        brainevent = BrainEvent()
 
-
-    brainevent = BrainEvent()
+else:
+    import brainevent

brainstate/mixin.py

@@ -132,6 +132,7 @@ class AlignPost(Mixin):
         raise NotImplementedError
 
 
+
 class BindCondData(Mixin):
     """Bind temporary conductance data.
 
@@ -147,7 +148,6 @@ class BindCondData(Mixin):
 
 
 class UpdateReturn(Mixin):
-
     def update_return(self) -> PyTree:
         """
         The update function return of the model.
@@ -157,19 +157,6 @@ class UpdateReturn(Mixin):
         """
         raise NotImplementedError(f'Must implement the "{self.update_return.__name__}()" function.')
 
-    def update_return_info(self) -> PyTree:
-        """
-        The update return information of the model.
-
-        It should be a pytree, with each element as a ``jax.Array``.
-
-        .. note::
-           Should not include the batch axis and batch in_size.
-           These information will be inferred from the ``mode`` attribute.
-
-        """
-        raise NotImplementedError(f'Must implement the "{self.update_return_info.__name__}()" function.')
-
 
 class _MetaUnionType(type):
     def __new__(cls, name, bases, dct):

brainstate/nn/__init__.py

@@ -19,46 +19,110 @@ from ._collective_ops import *
 from ._collective_ops import __all__ as collective_ops_all
 from ._common import *
 from ._common import __all__ as common_all
-from ._dyn_impl import *
-from ._dyn_impl import __all__ as dyn_impl_all
+from ._conv import *
+from ._conv import __all__ as conv_all
+from ._delay import *
+from ._delay import __all__ as state_delay_all
+from ._dropout import *
+from ._dropout import __all__ as dropout_all
 from ._dynamics import *
-from ._dynamics import __all__ as dynamics_all
+from ._dynamics import __all__ as dyn_all
 from ._elementwise import *
 from ._elementwise import __all__ as elementwise_all
-from ._event import *
-from ._event import __all__ as _event_all
+from ._embedding import *
+from ._embedding import __all__ as embed_all
 from ._exp_euler import *
 from ._exp_euler import __all__ as exp_euler_all
-from ._interaction import *
-from ._interaction import __all__ as interaction_all
+from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
+from ._inputs import *
+from ._inputs import __all__ as inputs_all
+from ._linear import *
+from ._linear import __all__ as linear_all
+from ._linear_mv import EventLinear
+from ._ltp import *
+from ._ltp import __all__ as ltp_all
 from ._module import *
 from ._module import __all__ as module_all
+from ._neuron import *
+from ._neuron import __all__ as dyn_neuron_all
+from ._normalizations import *
+from ._normalizations import __all__ as normalizations_all
+from ._poolings import *
+from ._poolings import __all__ as poolings_all
+from ._projection import *
+from ._projection import __all__ as projection_all
+from ._rate_rnns import *
+from ._rate_rnns import __all__ as rate_rnns
+from ._readout import *
+from ._readout import __all__ as readout_all
+from ._stp import *
+from ._stp import __all__ as stp_all
+from ._synapse import *
+from ._synapse import __all__ as dyn_synapse_all
+from ._synaptic_projection import *
+from ._synaptic_projection import __all__ as _syn_proj_all
+from ._synouts import *
+from ._synouts import __all__ as synouts_all
 from ._utils import *
 from ._utils import __all__ as utils_all
 
 __all__ = (
-    ['metrics']
+    [
+        'metrics',
+        'EventLinear',
+        'EventFixedProb',
+        'EventFixedNumConn',
+    ]
     + collective_ops_all
     + common_all
-    + dyn_impl_all
-    + dynamics_all
     + elementwise_all
     + module_all
     + exp_euler_all
-    + interaction_all
     + utils_all
-    + _event_all
+    + dyn_all
+    + projection_all
+    + state_delay_all
+    + synouts_all
+    + conv_all
+    + linear_all
+    + normalizations_all
+    + poolings_all
+    + embed_all
+    + dropout_all
+    + elementwise_all
+    + dyn_neuron_all
+    + dyn_synapse_all
+    + inputs_all
+    + rate_rnns
+    + readout_all
+    + stp_all
+    + ltp_all
+    + _syn_proj_all
 )
 
 del (
     collective_ops_all,
     common_all,
-    dyn_impl_all,
-    dynamics_all,
-    elementwise_all,
     module_all,
     exp_euler_all,
-    interaction_all,
     utils_all,
-    _event_all,
+    dyn_all,
+    projection_all,
+    state_delay_all,
+    synouts_all,
+    conv_all,
+    linear_all,
+    normalizations_all,
+    poolings_all,
+    embed_all,
+    dropout_all,
+    elementwise_all,
+    dyn_neuron_all,
+    dyn_synapse_all,
+    inputs_all,
+    readout_all,
+    rate_rnns,
+    stp_all,
+    ltp_all,
+    _syn_proj_all,
 )

brainstate/nn/{_interaction/_conv.py → _conv.py}

@@ -23,8 +23,8 @@ import jax.numpy as jnp
 
 from brainstate import init, functional
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
+from ._module import Module
 
 T = TypeVar('T')
 

brainstate/nn/{_dynamics/_state_delay.py → _delay.py}

@@ -27,9 +27,9 @@ from brainstate import environ
 from brainstate._state import ShortTermState, State
 from brainstate.compile import jit_error_if
 from brainstate.graph import Node
-from brainstate.nn._collective_ops import call_order
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, PyTree
+from ._collective_ops import call_order
+from ._module import Module
 
 __all__ = [
     'Delay', 'DelayAccess', 'StateWithDelay',
@@ -135,6 +135,7 @@ class Delay(Module):
         entries: Optional[Dict] = None,  # delay access entry
         delay_method: Optional[str] = _DELAY_ROTATE,  # delay method
         interp_method: str = _INTERP_LINEAR,  # interpolation method
+        take_aware_unit: bool = False
     ):
         # target information
         self.target_info = jax.tree.map(lambda a: jax.ShapeDtypeStruct(a.shape, a.dtype), target_info)
@@ -170,6 +171,9 @@ class Delay(Module):
             for entry, delay_time in entries.items():
                 self.register_entry(entry, delay_time)
 
+        self.take_aware_unit = take_aware_unit
+        self._unit = None
+
     @property
     def history(self):
         return self._history

brainstate/nn/{_elementwise/_dropout.py → _dropout.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import random, environ, init
 from brainstate._state import ShortTermState
-from brainstate.nn._module import ElementWiseBlock
 from brainstate.typing import Size
+from ._module import ElementWiseBlock
 
 __all__ = [
     'DropoutFixed', 'Dropout', 'Dropout1d', 'Dropout2d', 'Dropout3d',

brainstate/nn/{_dynamics/_dynamics_base.py → _dynamics.py}

@@ -41,13 +41,14 @@ import numpy as np
 from brainstate import environ
 from brainstate._state import State
 from brainstate.graph import Node
-from brainstate.mixin import ParamDescriber
-from brainstate.nn._module import Module
+from brainstate.mixin import ParamDescriber, UpdateReturn
 from brainstate.typing import Size, ArrayLike
-from ._state_delay import StateWithDelay, Delay
+from ._delay import StateWithDelay, Delay
+from ._module import Module
 
 __all__ = [
-    'DynamicsGroup', 'Projection', 'Dynamics', 'Prefetch',
+    'DynamicsGroup', 'Projection', 'Dynamics',
+    'Prefetch', 'PrefetchDelay', 'PrefetchDelayAt', 'OutputDelayAt',
 ]
 
 T = TypeVar('T')
@@ -99,7 +100,7 @@ class Projection(Module):
             raise ValueError('Do not implement the update() function.')
 
 
-class Dynamics(Module):
+class Dynamics(Module, UpdateReturn):
     """
     Base class for implementing neural dynamics models in BrainState.
 
@@ -821,6 +822,41 @@ class Dynamics(Module):
         else:
             raise TypeError(f'The input {dyn} should be an instance of {Dynamics} or a delayed initializer.')
 
+    def prefetch_delay(self, state: str, delay: Optional[ArrayLike] = None) -> 'PrefetchDelayAt':
+        """
+        Create a reference to a delayed state or variable in the module.
+
+        This method simplifies the process of accessing a delayed version of a state or variable
+        within the module. It first creates a prefetch reference to the specified state,
+        then specifies the delay time for accessing this state.
+
+        Args:
+            state (str): The name of the state or variable to reference.
+            delay (Optional[ArrayLike]): The amount of time to delay the variable access,
+                typically in time units (e.g., milliseconds). Defaults to None.
+
+        Returns:
+            PrefetchDelayAt: An object that provides access to the variable at the specified delay time.
+        """
+        return self.prefetch(state).delay.at(delay)
+
+    def output_delay(self, delay: Optional[ArrayLike] = None) -> 'OutputDelayAt':
+        """
+        Create a reference to the delayed output of the module.
+
+        This method simplifies the process of accessing a delayed version of the module's output.
+        It instantiates an `OutputDelayAt` object, which can be used to retrieve the output value
+        at the specified delay time.
+
+        Args:
+            delay (Optional[ArrayLike]): The amount of time to delay the output access,
+                typically in time units (e.g., milliseconds). Defaults to None.
+
+        Returns:
+            OutputDelayAt: An object that provides access to the module's output at the specified delay time.
+        """
+        return OutputDelayAt(self, delay)
+
 
 class Prefetch(Node):
     """
@@ -885,6 +921,7 @@ class Prefetch(Node):
             An object that provides access to delayed versions of the prefetched item.
         """
         return PrefetchDelay(self.module, self.item)
+        # return PrefetchDelayAt(self.module, self.item, time)
 
     def __call__(self, *args, **kwargs):
         """
@@ -1007,7 +1044,7 @@ class PrefetchDelayAt(Node):
         self,
         module: Dynamics,
         item: str,
-        time: ArrayLike
+        time: ArrayLike = None,
     ):
         """
         Initialize a PrefetchDelayAt object.
@@ -1026,14 +1063,16 @@ class PrefetchDelayAt(Node):
         self.module = module
         self.item = item
         self.time = time
-        self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
 
-        # register the delay
-        key = _get_delay_key(item)
-        if not module._has_after_update(key):
-            module._add_after_update(key, not_receive_update_output(StateWithDelay(module, item)))
-        self.state_delay: StateWithDelay = module._get_after_update(key)
-        self.state_delay.register_delay(time)
+        if time is not None:
+            self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
+
+            # register the delay
+            key = _get_prefetch_delay_key(item)
+            if not module._has_after_update(key):
+                module._add_after_update(key, not_receive_update_output(StateWithDelay(module, item)))
+            self.state_delay: StateWithDelay = module._get_after_update(key)
+            self.state_delay.register_delay(time)
 
     def __call__(self, *args, **kwargs):
         """
@@ -1044,12 +1083,94 @@ class PrefetchDelayAt(Node):
         Any
             The value of the state or variable at the specified delay time.
         """
-        # return self.state_delay.retrieve_at_time(self.time)
-        return self.state_delay.retrieve_at_step(self.step)
+        if self.time is None:
+            return _get_prefetch_item(self).value
+        else:
+            return self.state_delay.retrieve_at_step(self.step)
+
+
+class OutputDelayAt(Node):
+    """
+    Provides access to a specific delayed state or variable value at the specific time.
+
+    This class represents the final step in the prefetch delay chain, providing
+    actual access to state values at a specific delay time. It converts the
+    specified time delay into steps and registers the delay with the appropriate
+    StateWithDelay handler.
+
+    Parameters
+    ----------
+    module : Dynamics
+        The dynamics module that contains the referenced state or variable.
+    item : str
+        The name of the state or variable to access with delay.
+    time : ArrayLike
+        The amount of time to delay access by, typically in time units (e.g., milliseconds).
+
+    Examples
+    --------
+    >>> import brainstate
+    >>> import brainunit as u
+    >>> neuron = brainstate.nn.LIF(10)
+    >>> # Create a reference to voltage delayed by 5ms
+    >>> delayed_v = PrefetchDelayAt(neuron, 'V', 5.0 * u.ms)
+    >>> # Get the delayed value
+    >>> v_value = delayed_v()
+    """
+
+    def __init__(
+        self,
+        module: Dynamics,
+        time: Optional[ArrayLike] = None,
+    ):
+        """
+        Initialize a PrefetchDelayAt object.
+
+        Parameters
+        ----------
+        module : AlignPre, Module
+            The dynamics module that contains the referenced state or variable.
+        time : ArrayLike
+            The amount of time to delay access by, typically in time units.
+        """
+        super().__init__()
+        assert isinstance(module, UpdateReturn), 'The module should implement the `update_return` method.'
+        assert isinstance(module, Module), 'The module should be an instance of Module.'
+        self.module = module
+        self.time = time
+        if time is not None:
+            self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
+
+            # register the delay
+            key = _get_output_delay_key()
+            if not module._has_after_update(key):
+                # TODO: unit processing
+                delay = Delay(module.update_return(), time)
+                module._add_after_update(key, receive_update_output(delay))
+            self.out_delay: Delay = module._get_after_update(key)
+            self.out_delay.register_delay(time)
+
+    def __call__(self, *args, **kwargs):
+        """
+        Retrieve the value of the state at the specified delay time.
+
+        Returns
+        -------
+        Any
+            The value of the state or variable at the specified delay time.
+        """
+        if self.time is None:
+            return self.module.update_return()
+        else:
+            return self.out_delay.retrieve_at_step(self.step)
+
+
+def _get_prefetch_delay_key(item) -> str:
+    return f'{item}-prefetch-delay'
 
 
-def _get_delay_key(item) -> str:
-    return f'{item}-delay'
+def _get_output_delay_key() -> str:
+    return f'output-delay'
 
 
 def _get_prefetch_item(target: Union[Prefetch, PrefetchDelayAt]) -> Any:
@@ -1064,7 +1185,7 @@ def _get_prefetch_item_delay(target: Union[Prefetch, PrefetchDelay, PrefetchDela
         f'The target module should be an instance '
         f'of Dynamics. But got {target.module}.'
     )
-    delay = target.module._get_after_update(_get_delay_key(target.item))
+    delay = target.module._get_after_update(_get_prefetch_delay_key(target.item))
     if not isinstance(delay, StateWithDelay):
         raise TypeError(f'The prefetch target should be a {StateWithDelay.__name__} when accessing '
                         f'its delay. But got {delay}.')

brainstate/nn/{_elementwise/_elementwise.py → _elementwise.py}

@@ -23,8 +23,8 @@ import jax.typing
 
 from brainstate import random, functional as F
 from brainstate._state import ParamState
-from brainstate.nn._module import ElementWiseBlock
 from brainstate.typing import ArrayLike
+from ._module import ElementWiseBlock
 
 __all__ = [
     # activation functions

brainstate/nn/{_interaction/_embedding.py → _embedding.py}

@@ -17,8 +17,8 @@ from typing import Optional, Callable, Union
 
 from brainstate import init
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
+from ._module import Module
 
 __all__ = [
     'Embedding',

brainstate/nn/{_event/_fixedprob_mv.py → _fixedprob_mv.py}

@@ -25,8 +25,8 @@ from brainstate import random, augment, environ, init
 from brainstate._compatible_import import brainevent
 from brainstate._state import ParamState
 from brainstate.compile import for_loop
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
+from ._module import Module
 
 __all__ = [
     'EventFixedNumConn',

brainstate/nn/{_dyn_impl/_inputs.py → _inputs.py}

@@ -20,12 +20,11 @@ import jax
 import numpy as np
 
 from brainstate import environ, init, random
-from brainstate._state import ShortTermState
-from brainstate._state import State, maybe_state
+from brainstate._state import ShortTermState, State, maybe_state
 from brainstate.compile import while_loop
-from brainstate.nn._dynamics._dynamics_base import Dynamics, Prefetch
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size, DTypeLike
+from ._dynamics import Dynamics, Prefetch
+from ._module import Module
 
 __all__ = [
     'SpikeTime',
@@ -134,7 +133,7 @@ class PoissonSpike(Dynamics):
         self.freqs = init.param(freqs, self.varshape, allow_none=False)
 
     def update(self):
-        spikes = random.rand(self.varshape) <= (self.freqs * environ.get_dt())
+        spikes = random.rand(*self.varshape) <= (self.freqs * environ.get_dt())
         spikes = u.math.asarray(spikes, dtype=self.spk_type)
         return spikes
 

brainstate/nn/{_interaction/_linear.py → _linear.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import init, functional
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size
+from ._module import Module
 
 __all__ = [
     'Linear',
@@ -350,10 +350,7 @@ class OneToOne(Module):
         self.weight = param_type(param)
 
     def update(self, pre_val):
-        pre_val, pre_unit = u.get_mantissa(pre_val), u.get_unit(pre_val)
-        w_val, w_unit = u.get_mantissa(self.weight.value['weight']), u.get_unit(self.weight.value['weight'])
-        post_val = pre_val * w_val
-        post_val = u.maybe_decimal(u.Quantity(post_val, unit=w_unit * pre_unit))
+        post_val = pre_val * self.weight.value['weight']
         if 'bias' in self.weight.value:
             post_val = post_val + self.weight.value['bias']
         return post_val

brainstate/nn/{_event/_linear_mv.py → _linear_mv.py}

@@ -21,8 +21,8 @@ import jax
 from brainstate import init
 from brainstate._compatible_import import brainevent
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
+from ._module import Module
 
 __all__ = [
     'EventLinear',

brainstate/nn/{_event/__init__.py → _ltp.py}

@@ -16,11 +16,13 @@
 # -*- coding: utf-8 -*-
 
 
-from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
-from ._linear_mv import EventLinear
+from ._synapse import Synapse
 
 __all__ = [
-    'EventLinear',
-    'EventFixedProb',
-    'EventFixedNumConn',
+    'LongTermPlasticity',
 ]
+
+
+class LongTermPlasticity(Synapse):
+    pass
+

brainstate/nn/{_dyn_impl/_dynamics_neuron.py → _neuron.py}

@@ -22,9 +22,9 @@ import jax
 
 from brainstate import init, surrogate, environ
 from brainstate._state import HiddenState, ShortTermState
-from brainstate.nn._dynamics._dynamics_base import Dynamics
-from brainstate.nn._exp_euler import exp_euler_step
 from brainstate.typing import ArrayLike, Size
+from ._dynamics import Dynamics
+from ._exp_euler import exp_euler_step
 
 __all__ = [
     'Neuron', 'IF', 'LIF', 'LIFRef', 'ALIF',

brainstate/nn/{_interaction/_normalizations.py → _normalizations.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import environ, init
 from brainstate._state import ParamState, BatchState
-from brainstate.nn._module import Module
 from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
+from ._module import Module
 
 __all__ = [
     'BatchNorm0d',

brainstate/nn/{_interaction/_poolings.py → _poolings.py}

@@ -25,8 +25,8 @@ import jax.numpy as jnp
 import numpy as np
 
 from brainstate import environ
-from brainstate.nn._module import Module
 from brainstate.typing import Size
+from ._module import Module
 
 __all__ = [
     'Flatten', 'Unflatten',

brainstate/nn/{_interaction/_poolings_test.py → _poolings_test.py}

@@ -103,7 +103,7 @@ class TestPool(parameterized.TestCase):
         for target_size in [10, 9, 8, 7, 6]
     )
     def test_adaptive_pool1d(self, target_size):
-        from brainstate.nn._interaction._poolings import _adaptive_pool1d
+        from brainstate.nn._poolings import _adaptive_pool1d
 
         arr = brainstate.random.rand(100)
         op = jax.numpy.mean