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

brainstate/nn/__init__.py

@@ -19,46 +19,113 @@ 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 import *
+from._fixedprob import __all__ as fixedprob_all
+from ._inputs import *
+from ._inputs import __all__ as inputs_all
+from ._linear import *
+from ._linear import __all__ as linear_all
+from ._linear_mv import *
+from ._linear_mv import __all__ as linear_mv_all
+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',
+    ]
     + 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
+    + fixedprob_all
+    + linear_mv_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,
+    fixedprob_all,
+    linear_mv_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
@@ -326,7 +330,14 @@ class Delay(Module):
         indices = (delay_idx,) + indices
 
         # the delay data
-        return jax.tree.map(lambda a: a[indices], self.history.value)
+        if self._unit is None:
+            return jax.tree.map(lambda a: a[indices], self.history.value)
+        else:
+            return jax.tree.map(
+                lambda hist, unit: u.maybe_decimal(hist[indices] * unit),
+                self.history.value,
+                self._unit
+            )
 
     def retrieve_at_time(self, delay_time, *indices) -> PyTree:
         """
@@ -389,6 +400,9 @@ class Delay(Module):
         """
         assert self.history is not None, 'The delay history is not initialized.'
 
+        if self.take_aware_unit and self._unit is None:
+            self._unit = jax.tree.map(lambda x: u.get_unit(x), current, is_leaf=u.math.is_quantity)
+
         # update the delay data at the rotation index
         if self.delay_method == _DELAY_ROTATE:
             i = environ.get(environ.I)
@@ -415,6 +429,8 @@ class Delay(Module):
             raise ValueError(f'Unknown updating method "{self.delay_method}"')
 
 
+
+
 class StateWithDelay(Delay):
     """
     A ``State`` type that defines the state in a differential equation.

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',
@@ -409,7 +409,8 @@ class DropoutFixed(ElementWiseBlock):
         self.out_size = in_size
 
     def init_state(self, batch_size=None, **kwargs):
-        self.mask = ShortTermState(init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size))
+        if self.prob < 1.:
+            self.mask = ShortTermState(init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size))
 
     def update(self, x):
         dtype = u.math.get_dtype(x)
@@ -418,8 +419,8 @@ class DropoutFixed(ElementWiseBlock):
             if self.mask.value.shape != x.shape:
                 raise ValueError(f"Input shape {x.shape} does not match the mask shape {self.mask.value.shape}. "
                                  f"Please call `init_state()` method first.")
-            return jnp.where(self.mask.value,
-                             jnp.asarray(x / self.prob, dtype=dtype),
-                             jnp.asarray(0., dtype=dtype))
+            return u.math.where(self.mask.value,
+                                u.math.asarray(x / self.prob, dtype=dtype),
+                                u.math.asarray(0., dtype=dtype) * u.get_unit(x))
         else:
             return x

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

@@ -36,18 +36,20 @@ For handling the delays:
 from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING
 
 import brainunit as u
+import jax
 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.typing import Size, ArrayLike
-from ._state_delay import StateWithDelay, Delay
+from brainstate.mixin import ParamDescriber, UpdateReturn
+from brainstate.typing import Size, ArrayLike, PyTree
+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 +101,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.
 
@@ -810,17 +812,64 @@ class Dynamics(Module):
         >>> n1.align_pre(brainstate.nn.Expon.desc(n1.varshape))  # n2 will run after n1
         """
         if isinstance(dyn, Dynamics):
-            self._add_after_update(dyn.name, dyn)
+            self._add_after_update(id(dyn), dyn)
             return dyn
         elif isinstance(dyn, ParamDescriber):
             if not issubclass(dyn.cls, Dynamics):
                 raise TypeError(f'The input {dyn} should be an instance of {Dynamics}.')
             if not self._has_after_update(dyn.identifier):
-                self._add_after_update(dyn.identifier, dyn())
+                self._add_after_update(
+                    dyn.identifier,
+                    dyn() if ('in_size' in dyn.kwargs or len(dyn.args) > 0) else dyn(in_size=self.varshape)
+                )
             return self._get_after_update(dyn.identifier)
         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,
+        variable_like: PyTree = 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.
+            variable_like:
+
+        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 +934,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 +1057,7 @@ class PrefetchDelayAt(Node):
         self,
         module: Dynamics,
         item: str,
-        time: ArrayLike
+        time: ArrayLike = None,
     ):
         """
         Initialize a PrefetchDelayAt object.
@@ -1026,14 +1076,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 +1096,76 @@ 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.
+    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_spike = OutputDelayAt(neuron, 5.0 * u.ms)
+    >>> # Get the delayed value
+    >>> v_value = delayed_spike()
+    """
+
+    def __init__(
+        self,
+        module: Dynamics,
+        time: Optional[ArrayLike] = None,
+        variable_like: Optional[PyTree] = None,
+    ):
+        super().__init__()
+        assert isinstance(module, Dynamics), 'The module should be an instance of Dynamics.'
+        self.module = module
+        dt = environ.get_dt()
+        if time is None:
+            time = u.math.zeros_like(dt)
+        self.time = time
+        self.step = u.math.asarray(time / dt, dtype=environ.ditype())
+
+        # register the delay
+        key = _get_output_delay_key()
+        if not module._has_after_update(key):
+            delay = Delay(
+                jax.ShapeDtypeStruct(module.out_size, dtype=environ.dftype()),
+                time,
+                take_aware_unit=True
+            )
+            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):
+        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 +1180,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',