brainstate 0.1.9__py2.py3-none-any.whl → 0.1.10__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.9"
+__version__ = "0.1.10"
 
 from . import augment
 from . import compile

brainstate/environ.py

@@ -402,13 +402,13 @@ def _get_float(precision: int):
 
 @functools.lru_cache()
 def _get_complex(precision: int):
-    if precision == [64, '64']:
+    if precision in [64, '64']:
         return np.complex128
-    elif precision == [32, '32']:
+    elif precision in [32, '32']:
         return np.complex64
     elif precision in [16, '16', 'bf16']:
         return np.complex64
-    elif precision == [8, '8']:
+    elif precision in [8, '8']:
         return np.complex64
     else:
         raise ValueError(f'Unsupported precision: {precision}')

brainstate/nn/_delay.py

@@ -472,15 +472,94 @@ class Delay(Module):
 
 class StateWithDelay(Delay):
     """
-    A ``State`` type that defines the state in a differential equation.
+    Delayed history buffer bound to a module state.
+
+    StateWithDelay is a specialized :py:class:`~.Delay` that attaches to a
+    concrete :py:class:`~brainstate._state.State` living on a target module
+    (for example a membrane potential ``V`` on a neuron). It automatically
+    maintains a rolling history of that state and exposes convenient helpers to
+    retrieve the value at a given delay either by step or by time.
+
+    In normal usage you rarely instantiate this class directly. It is created
+    implicitly when using the prefetch-delay helpers on a Dynamics module, e.g.:
+
+    - ``module.prefetch('V').delay.at(5.0 * u.ms)``
+    - ``module.prefetch_delay('V', 5.0 * u.ms)``
+
+    Both will construct a StateWithDelay bound to ``module.V`` under the hood
+    and register the requested delay, so you can retrieve the delayed value
+    inside your update rules.
+
+    Parameters
+    ----------
+    target : :py:class:`~brainstate.graph.Node`
+        The module object that owns the state to track.
+    item : str
+        The attribute name of the target state on ``target`` (must be a
+        :py:class:`~brainstate._state.State`).
+    init : Callable, optional
+        Optional initializer used to fill the history buffer before ``t0``
+        when delays request values from the past that hasn't been simulated yet.
+        The callable receives ``(shape, dtype)`` and must return an array.
+        If not provided, zeros are used. You may also pass a scalar/array
+        literal via the underlying Delay API when constructing manually.
+    delay_method : {"rotation", "concat"}, default "rotation"
+        Internal buffering strategy (inherits behavior from :py:class:`~.Delay`).
+        "rotation" keeps a ring buffer; "concat" shifts by concatenation.
+
+    Attributes
+    ----------
+    state : :py:class:`~brainstate._state.State`
+        The concrete state object being tracked.
+    history : :py:class:`~brainstate._state.ShortTermState`
+        Rolling time axis buffer with shape ``[length, *state.shape]``.
+    max_time : float
+        Maximum time span currently supported by the buffer.
+    max_length : int
+        Buffer length in steps (``ceil(max_time/dt)+1``).
+
+    Notes
+    -----
+    - This class inherits all retrieval utilities from :py:class:`~.Delay`:
+      use :py:meth:`retrieve_at_step` when you know the integer delay steps,
+      or :py:meth:`retrieve_at_time` for continuous-time queries with optional
+      linear/round interpolation.
+    - It is registered as an "after-update" hook on the owning Dynamics so the
+      buffer is updated automatically after each simulation step.
+
+    Examples
+    --------
+    Access a neuron's membrane potential 5 ms in the past:
+
+    >>> import brainunit as u
+    >>> import brainstate as bst
+    >>> lif = bst.nn.LIF(100)
+    >>> # Create a delayed accessor to V(t-5ms)
+    >>> v_delay = lif.prefetch_delay('V', 5.0 * u.ms)
+    >>> # Inside another module's update you can read the delayed value
+    >>> v_t_minus_5ms = v_delay()
+
+    Register multiple delay taps and index-specific delays:
+
+    >>> # Under the hood, a StateWithDelay is created and you can register
+    >>> # additional taps (in steps or time) via its Delay interface
+    >>> _ = lif.prefetch('V').delay.at(2.0 * u.ms)   # additional delay
+    >>> # Direct access to buffer by steps (advanced)
+    >>> # lif._get_after_update('V-prefetch-delay').retrieve_at_step(3)
     """
 
     __module__ = 'brainstate.nn'
 
     state: State  # state
 
-    def __init__(self, target: Node, item: str, init: Callable = None):
-        super().__init__(None, init=init)
+    def __init__(
+        self,
+        target: Node,
+        item: str,
+        init: Callable = None,
+        delay_method: Optional[str] = _DELAY_ROTATE,
+    ):
+        super().__init__(None, init=init, delay_method=delay_method)
 
         self._target = target
         self._target_term = item

{brainstate-0.1.9.dist-info → brainstate-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: brainstate
-Version: 0.1.9
+Version: 0.1.10
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.9.dist-info → brainstate-0.1.10.dist-info}/RECORD

@@ -1,9 +1,9 @@
-brainstate/__init__.py,sha256=4vCHesk-UDMgifDdr5hm7veKU8Avg94-DC0OLXrQn8g,1450
+brainstate/__init__.py,sha256=lG0hJsHre06CcJMm_Em25hkKKFH9J8ZqnyGPq3iixZI,1451
 brainstate/_compatible_import.py,sha256=kCJy0OIr2VAJHqyk2bdAZ1s2VAFKeEgOpfFosy9wi7A,4482
 brainstate/_state.py,sha256=iZeLyjaCG39yXHh7y5LkOC2_MwkLgEnnKKr6puHkeME,56966
 brainstate/_state_test.py,sha256=UqbXHhJNAHLehvVMnxnOUfWT1PVRWFawOvoyjJvGIRo,1586
 brainstate/_utils.py,sha256=44Yy4cBlsSfZln7CEF8XfNCsZiIIDPPmC9LT17OOi5s,1610
-brainstate/environ.py,sha256=yRIevZvU92sYz6XTqmDLAjCXCnrHW4_FwWKRPNygkf8,17230
+brainstate/environ.py,sha256=mVIb_4KUP7s26tBMIuEfoo47Q6lcnY_XlCIE5rRO5m8,17230
 brainstate/environ_test.py,sha256=KYnPRv-ZgNDrvoTyNGuyO5ElLBoSFR4AN-B1hi2tp4k,1983
 brainstate/mixin.py,sha256=88obor0L9qhoj4Kt_Q-dID9ZyPO-P7pasgbbykdWeoA,10651
 brainstate/mixin_test.py,sha256=xhKEQMvpHQE74o_hQ88EID-mL_tSF7eStCSCVUfY5Tk,3463
@@ -57,7 +57,7 @@ brainstate/nn/_collective_ops_test.py,sha256=bZlWQDldsdaX89exbhnXPTP-sONrk6riii3
 brainstate/nn/_common.py,sha256=N3AaSqRLDOtdynGE_xO03EJXvyh55G6O1q1ZrHKkhf8,6253
 brainstate/nn/_conv.py,sha256=r77NK9xAUy4l6bh680wAlV69FzzQzLE56gcxlLKn5as,17947
 brainstate/nn/_conv_test.py,sha256=Ekkra0Inn0aiBTJcF7orrrlpWXWs1Vv0xjcJxQ9GiDY,8624
-brainstate/nn/_delay.py,sha256=0B2t6--9I3XOTjdzufAjEYhgZg8zwzduqmmvKA_HpBk,19590
+brainstate/nn/_delay.py,sha256=jMCvnYrivXJwAmZQBfLxM29wSlkgmW2FudXGT_yt8kk,23051
 brainstate/nn/_delay_test.py,sha256=p0YbiqpqRe_ulhFr9qDmMt2qihXg5En9O_hWxAlHFVc,10158
 brainstate/nn/_dropout.py,sha256=xiKeTsLdq5a3kCmhjf9XzuTc2VMjF6NEqJFLUqqKfKw,17357
 brainstate/nn/_dropout_test.py,sha256=J0Ju3weOWWk6ZYkG7_3At8Kq-xnE5WK8_svfZBJAq04,4346
@@ -123,8 +123,8 @@ brainstate/util/pretty_repr.py,sha256=DVDwNRfXFrYe_j2HPJ9-1rur5tJ_YXMduImUOwjf7c
 brainstate/util/pretty_table.py,sha256=qiEA3lSd9l4-Evia2KSu_KSw-F8p2BmyB7XepLpDVlk,105163
 brainstate/util/scaling.py,sha256=Y4kYCr3dDmn808VUCY3i_L1XFLrUSOkRJykJukohCBA,7266
 brainstate/util/struct.py,sha256=RBvB_LZyAWt6-3SHJnUfS3pE7D4pcEdrYEPbjw_fGfU,17005
-brainstate-0.1.9.dist-info/licenses/LICENSE,sha256=RJ40fox7u2in2H8wvIS5DsPGlNHaA7JI024thFUlaZE,11348
-brainstate-0.1.9.dist-info/METADATA,sha256=L-hwNoEacz-zOcfjK6SYurldwS31HZjin3KcFCeZmRo,3625
-brainstate-0.1.9.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
-brainstate-0.1.9.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.9.dist-info/RECORD,,
+brainstate-0.1.10.dist-info/licenses/LICENSE,sha256=RJ40fox7u2in2H8wvIS5DsPGlNHaA7JI024thFUlaZE,11348
+brainstate-0.1.10.dist-info/METADATA,sha256=uZXs-VisW59P3ZC1GALwziVeF-XctLimfSfEbsIvvvI,3626
+brainstate-0.1.10.dist-info/WHEEL,sha256=JNWh1Fm1UdwIQV075glCn4MVuCRs0sotJIq-J6rbxCU,109
+brainstate-0.1.10.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.10.dist-info/RECORD,,