brainstate 0.1.0.post20241220__py2.py3-none-any.whl → 0.1.0.post20250101__py2.py3-none-any.whl

brainstate/__init__.py

@@ -48,19 +48,5 @@ __all__ = (
     _state_all
 )
 
-# ----------------------- #
-# deprecations
-# ----------------------- #
-
-from ._utils import deprecation_getattr
-
-transform._deprecations = dict()
-for key in compile.__all__:
-    transform._deprecations[key] = (f'brainstate.transform.{key}', f'brainstate.compile.{key}', getattr(compile, key))
-for key in augment.__all__:
-    transform._deprecations[key] = (f'brainstate.transform.{key}', f'brainstate.augment.{key}', getattr(augment, key))
-transform.__getattr__ = deprecation_getattr('brainstate.transform', transform._deprecations)
-del deprecation_getattr
-
 # ----------------------- #
 del _state_all

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -41,7 +41,8 @@ class Synapse(Dynamics):
 
 
 class Expon(Synapse, AlignPost):
-    r"""Exponential decay synapse model.
+    r"""
+    Exponential decay synapse model.
 
     Args:
       tau: float. The time constant of decay. [ms]
@@ -75,6 +76,89 @@ class Expon(Synapse, AlignPost):
         return self.g.value
 
 
+class DualExpon(Synapse, AlignPost):
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        tau_decay: ArrayLike = 10.0 * u.ms,
+        tau_rise: ArrayLike = 1.0 * u.ms,
+        A: Optional[ArrayLike] = None,
+        g_initializer: ArrayLike = init.ZeroInit(unit=u.mS),
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau_decay = init.param(tau_decay, self.varshape)
+        self.tau_rise = init.param(tau_rise, self.varshape)
+        A = self._format_dual_exp_A(A)
+        self.a = (self.tau_decay - self.tau_rise) / self.tau_rise / self.tau_decay * A
+        self.g_initializer = g_initializer
+
+    def _format_dual_exp_A(self, A):
+        A = init.param(A, sizes=self.varshape, allow_none=True)
+        if A is None:
+            A = (
+                self.tau_decay / (self.tau_decay - self.tau_rise) *
+                u.math.float_power(self.tau_rise / self.tau_decay,
+                                   self.tau_rise / (self.tau_rise - self.tau_decay))
+            )
+        return A
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.g_rise = HiddenState(init.param(self.g_initializer, self.varshape, batch_size))
+        self.g_decay = HiddenState(init.param(self.g_initializer, self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.g_rise.value = init.param(self.g_initializer, self.varshape, batch_size)
+        self.g_decay.value = init.param(self.g_initializer, self.varshape, batch_size)
+
+    def update(self, x=None):
+        g_rise = exp_euler_step(lambda h: -h / self.tau_rise, self.g_rise.value)
+        g_decay = exp_euler_step(lambda g: -g / self.tau_decay, self.g_decay.value)
+        self.g_rise.value = self.sum_delta_inputs(g_rise)
+        self.g_decay.value = self.sum_delta_inputs(g_decay)
+        if x is not None:
+            self.g_rise.value += x
+            self.g_decay.value += x
+        return self.a * (self.g_decay.value - self.g_rise.value)
+
+
+class Alpha(Synapse):
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        tau: ArrayLike = 8.0 * u.ms,
+        g_initializer: ArrayLike = init.ZeroInit(unit=u.mS),
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau = init.param(tau, self.varshape)
+        self.g_initializer = g_initializer
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.g = HiddenState(init.param(self.g_initializer, self.varshape, batch_size))
+        self.h = HiddenState(init.param(self.g_initializer, self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.g.value = init.param(self.g_initializer, self.varshape, batch_size)
+        self.h.value = init.param(self.g_initializer, self.varshape, batch_size)
+
+    def update(self, x=None):
+        h = exp_euler_step(lambda h: -h / self.tau, self.h.value)
+        self.g.value = exp_euler_step(lambda g, h: -g / self.tau + h / self.tau, self.g.value, self.h.value)
+        self.h.value = self.sum_delta_inputs(h)
+        if x is not None:
+            self.h.value += x
+        return self.g.value
+
+
 class STP(Synapse):
     r"""Synaptic output with short-term plasticity.
 

brainstate/nn/_dynamics/_state_delay.py

@@ -51,11 +51,14 @@ def _get_delay(delay_time, delay_step):
             assert isinstance(delay_step, int), '"delay_step" should be an integer.'
             if delay_step == 0:
                 return 0., 0
-            delay_time = delay_step * environ.get_dt()
+            with jax.ensure_compile_time_eval():
+                delay_time = delay_step * environ.get_dt()
     else:
         assert delay_step is None, '"delay_step" should be None if "delay_time" is given.'
         # assert isinstance(delay_time, (int, float))
-        delay_step = math.ceil(delay_time / environ.get_dt())
+        with jax.ensure_compile_time_eval():
+            delay_step = delay_time / environ.get_dt()
+        delay_step = math.ceil(float(delay_step))
     return delay_time, delay_step
 
 

brainstate/transform.py

@@ -13,4 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
-# This module is going to be deleted in the future (near 2025-06).
+# alias for compilation and augmentation functions
+
+from .compile import *
+from .augment import *
+

{brainstate-0.1.0.post20241220.dist-info → brainstate-0.1.0.post20250101.dist-info}/METADATA

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

{brainstate-0.1.0.post20241220.dist-info → brainstate-0.1.0.post20250101.dist-info}/RECORD

@@ -1,4 +1,4 @@
-brainstate/__init__.py,sha256=r7C3eLTg8LEusoH6PGgBFFt4ZgbketYLoLA0lQhUCsE,2098
+brainstate/__init__.py,sha256=A-QKdOvSalsCMxgk80Iz6_xMiUin6con6JaONHfciSY,1526
 brainstate/_state.py,sha256=4aDpLyHGr1VlPXeLSfM3USQG5K4o7orF7IlaBdYrtfE,29098
 brainstate/_state_test.py,sha256=1boTp1w8DiCFLsPwNtlLrlIqGRpkasAmLid5bv2fgP4,2223
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
@@ -7,7 +7,7 @@ brainstate/environ_test.py,sha256=jXX3nR1CO74aow5YqfqSd73isj9MWgHQxrwSsEjTDY8,19
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
 brainstate/surrogate.py,sha256=YaY6RJ6kzpuPXWFjaWsxWt2MzJfdm5v_jeOR8V_jPoU,48369
-brainstate/transform.py,sha256=uryXL2vAD-oApUNvL4n6uzLPRwXXARCJ6lRwco_Gh_Y,772
+brainstate/transform.py,sha256=cxbymTlJ6uHvJWEEYXzFUkAySs_TbUTHakt0NQgWJ3s,808
 brainstate/typing.py,sha256=Qh-LBzm6oG4rSXv4V5qB8SNYcoOR7bASoK_iQxnlafk,10467
 brainstate/augment/__init__.py,sha256=BtXIBel7GbttmfBX6grxOxl0IiOJxLEa7qCGAXumamE,1286
 brainstate/augment/_autograd.py,sha256=o9ivoEY7BmtdM1XmzdMmeRXpj6Tvn5xNB8LSGp2HKC8,25238
@@ -81,7 +81,7 @@ brainstate/nn/metrics.py,sha256=iupHjSRTHYY-HmEPBC4tXWrZfF4zh1ek2NwSAA0gnwE,1473
 brainstate/nn/_dyn_impl/__init__.py,sha256=Oazar7h89dp1WA2Vx4Tj7gCBhxJKH4LAUEABkBEG7vU,1462
 brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=cTbIn41EPYG0h3ICzKBXxpgB6wwA2K8k5FAcf3Pa5N8,10927
 brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=Tfzrzu7udGrLJGnqItiLWe5WT0dgduvYOgzGCnaPJQg,6317
-brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=IHy6IsGjWpKZ8NLq4X7PaRwx3tpO2HRZNppCWM2fe4I,11862
+brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=41P3qJf8fi4J7dZdDnjChJF6lYJjFAOkgy9aE3FReY4,15247
 brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=t5i-HV0ii9sUNzWTEv04o26QVtQ-mCdMJcFq2MD755A,4981
 brainstate/nn/_dyn_impl/_inputs.py,sha256=pkcAVt_o5kQF_BGCTZZ-NUQpHgjlFHHPwtYC0fJkAA0,9099
 brainstate/nn/_dyn_impl/_projection_alignpost.py,sha256=PNC1Tzx_SF2DHAHeJCufXzO_Q4qLoBpWABI45B3GRuc,876
@@ -93,7 +93,7 @@ brainstate/nn/_dynamics/__init__.py,sha256=j1HSWu01wf5-KjSaNhBC9utVGDALOhUsFPrLP
 brainstate/nn/_dynamics/_dynamics_base.py,sha256=SOfrET1Z-JkS2GPo2bkQ-v40FLozIxT5A7ObwW_hQYo,21987
 brainstate/nn/_dynamics/_dynamics_base_test.py,sha256=gXMwENqqSvyZbMpLP0QtYndJ_h39dF5gIeiiSbMAjTk,2721
 brainstate/nn/_dynamics/_projection_base.py,sha256=jYe3WdBMgz2TJkcxPWEkyK7OA4IR1ChISd2GTfM6U2o,13528
-brainstate/nn/_dynamics/_state_delay.py,sha256=nZYGmVKmQvAQu-W4YOUFH1gnr-ZS3rg_GNkRhI8rQ-I,16761
+brainstate/nn/_dynamics/_state_delay.py,sha256=KfsXXoFg8hv9wZgY-p5_-cTGbY88rviJmCfnGiDgboo,16905
 brainstate/nn/_dynamics/_synouts.py,sha256=9TGAc-nVa50th7KKn4oKLbro-4W4rwxYvp-eu7ksAIE,4491
 brainstate/nn/_dynamics/_synouts_test.py,sha256=V_jDswRN4VvEXD-2yJO3VA1TALgX0HK6oPBQiUntOWc,2266
 brainstate/nn/_elementwise/__init__.py,sha256=PK8oq1K_EG2941AiUyLxCWoRdWvMO3yt8ZJbw3Lkhu8,935
@@ -137,8 +137,8 @@ brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
 brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20241220.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20241220.dist-info/METADATA,sha256=lfkUbD1vYx4bikkUolrL_pCp4zryfvGExXyWP_QH5tM,3533
-brainstate-0.1.0.post20241220.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20241220.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20241220.dist-info/RECORD,,
+brainstate-0.1.0.post20250101.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250101.dist-info/METADATA,sha256=QFsDpwFlj0QnoV3srIVjxTiWxJEmVROOlhLRM3u-B44,3533
+brainstate-0.1.0.post20250101.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250101.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250101.dist-info/RECORD,,