brainstate 0.1.0.post20250127__py2.py3-none-any.whl → 0.1.0.post20250129__py2.py3-none-any.whl

brainstate/nn/_dyn_impl/_inputs.py

@@ -216,9 +216,14 @@ def poisson_input(
     weight: u.Quantity,
     target: State,
     indices: Optional[Union[np.ndarray, jax.Array]] = None,
+    refractory: Optional[Union[jax.Array]] = None,
 ):
     """
     Poisson Input to the given :py:class:`brainstate.State`.
+
+    Args:
+        freq:
+        refractory: should be the same length with ``target``.
     """
     freq = maybe_state(freq)
     weight = maybe_state(weight)
@@ -291,7 +296,7 @@ def poisson_input(
         # )
 
         # update target variable
-        target.value = jax.tree.map(
+        data = jax.tree.map(
             lambda tar, x: tar + x * weight,
             target.value,
             inp,
@@ -358,9 +363,19 @@ def poisson_input(
         # )
 
         # update target variable
-        target.value = jax.tree.map(
+        data = jax.tree.map(
             lambda x, tar: tar.at[indices].add(x * weight),
             inp,
             tar_val,
             is_leaf=u.math.is_quantity
         )
+
+    if refractory is not None:
+        target.value = jax.tree.map(
+            lambda x, tar: u.math.where(refractory, tar, x),
+            data,
+            tar_val,
+            is_leaf=u.math.is_quantity
+        )
+    else:
+        target.value = data

{brainstate-0.1.0.post20250127.dist-info → brainstate-0.1.0.post20250129.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250127
+Version: 0.1.0.post20250129
 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.post20250127.dist-info → brainstate-0.1.0.post20250129.dist-info}/RECORD

@@ -65,7 +65,7 @@ brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=cTbIn41EPYG0h3ICzKBXxpgB6wwA2
 brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=Tfzrzu7udGrLJGnqItiLWe5WT0dgduvYOgzGCnaPJQg,6317
 brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=MsbPyaiDyjungyzuK2b3exRGaMpZgmsmmNHNLjgxQKw,15269
 brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=t5i-HV0ii9sUNzWTEv04o26QVtQ-mCdMJcFq2MD755A,4981
-brainstate/nn/_dyn_impl/_inputs.py,sha256=QOUpAb2YJOE78uAvIS8Ep6MFcQHV-V6uRwmYvk5p9bk,11385
+brainstate/nn/_dyn_impl/_inputs.py,sha256=72-UnT-hpG03EvSYx72ldDhbgZwmaoOYxxkANpX6xpo,11779
 brainstate/nn/_dyn_impl/_projection_alignpost.py,sha256=PNC1Tzx_SF2DHAHeJCufXzO_Q4qLoBpWABI45B3GRuc,876
 brainstate/nn/_dyn_impl/_rate_rnns.py,sha256=dz_yT_6hJVhKulcjIARbGtmMzZqISws96CtBc6o5GOo,14768
 brainstate/nn/_dyn_impl/_rate_rnns_test.py,sha256=gNgtr-a4ZiU1XF9wFG1HiJ9fLosfWchVR9Zn1x39xt4,2452
@@ -117,8 +117,8 @@ brainstate/util/_others.py,sha256=jsPZwP-v_5HRV-LB5F0NUsiqr04y8bmGIsu_JMyVcbQ,14
 brainstate/util/_pretty_repr.py,sha256=bDpU4gbkS4B8cXBkiN8kBQNmruxiJzDRF-eIqzyeYnM,5716
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=KMMHcshOM20gYhSahNzWLxsTt-Rt3AeX3Uz26-rP9vI,17619
-brainstate-0.1.0.post20250127.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250127.dist-info/METADATA,sha256=-j8gDuJN37nbSkijFC9TtG5UbflMZZWK-sftY7fkeps,3585
-brainstate-0.1.0.post20250127.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250127.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250127.dist-info/RECORD,,
+brainstate-0.1.0.post20250129.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250129.dist-info/METADATA,sha256=g93rl2oDT8uoaL4wbHVDFtC7Y5F3kuuebI64WBKUTc0,3585
+brainstate-0.1.0.post20250129.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250129.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250129.dist-info/RECORD,,