brainstate 0.1.0.post20250209__py2.py3-none-any.whl → 0.1.0.post20250211__py2.py3-none-any.whl

brainstate/_state.py

@@ -263,9 +263,8 @@ class State(Generic[A], PrettyReprTree):
         """
         The data and its value.
         """
-        self.check_if_deleted()
         record_state_value_read(self)
-        return self._value
+        return self._read_value()
 
     @value.setter
     def value(self, v) -> None:
@@ -275,7 +274,14 @@ class State(Generic[A], PrettyReprTree):
         Args:
           v: The value.
         """
-        self.write_value(v)
+        # NOTE: the following order is important
+
+        if isinstance(v, State):  # value checking
+            raise ValueError('Cannot set value to a State, ' 'use `copy_from` method instead')
+        self._check_value_tree(v)  # check the tree structure
+        record_state_value_write(self)  # record the value by the stack (>= level)
+        self._been_writen = True  # set the flag
+        self._write_value(v)  # write the value
 
     @property
     def stack_level(self):
@@ -297,17 +303,18 @@ class State(Generic[A], PrettyReprTree):
         """
         self._level = level
 
-    def write_value(self, v) -> None:
-        # value checking
-        if isinstance(v, State):
-            raise ValueError('Cannot set value to a State, ' 'use `copy_from` method instead')
-        self._check_value_tree(v)
-        # write the value by the stack (>= level)
-        record_state_value_write(self)
-        # set the value
+    def _read_value(self) -> PyTree[ArrayLike]:
+        """
+        The interface to customize the value reading.
+        """
+        self.check_if_deleted()
+        return self._value
+
+    def _write_value(self, v) -> None:
+        """
+        The interface to customize the value writing.
+        """
         self._value = v
-        # set flag
-        self._been_writen = True
 
     def restore_value(self, v) -> None:
         """
@@ -813,11 +820,7 @@ class TreefyState(Generic[A], PrettyReprTree):
             return 'value', v
 
         if k == '_name':
-            if self.name is None:
-                return None, None
-            else:
-                return 'name', v
-
+            return (None, None) if v is None else ('name', v)
         return k, v
 
     def replace(self, value: B) -> TreefyState[B]:

brainstate/compile/_jit.py

@@ -82,7 +82,6 @@ def _get_jitted_fun(
             return fun.fun(*args, **params)
 
         # compile the function and get the state trace
-        # print('Compiling ...')
         state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
         read_state_vals = state_trace.get_read_state_values(True)
 

brainstate/compile/_make_jaxpr.py

@@ -499,7 +499,7 @@ class StatefulFunction(object):
         state_vals, out = self.jaxpr_call([st.value for st in state_trace.states], *args, **kwargs)
         for st, written, val in zip(state_trace.states, state_trace.been_writen, state_vals):
             if written:
-                st.write_value(val)
+                st.value = val
             else:
                 st.restore_value(val)
         return out

brainstate/compile/_util.py

@@ -31,7 +31,7 @@ def write_back_state_values(
     assert len(state_trace.states) == len(state_trace.been_writen) == len(read_state_vals) == len(write_state_vals)
     for st, write, val_r, val_w in zip(state_trace.states, state_trace.been_writen, read_state_vals, write_state_vals):
         if write:
-            st.write_value(val_w)
+            st.value = val_w
         else:
             st.restore_value(val_r)
 

brainstate/graph/_graph_operation.py

@@ -609,7 +609,7 @@ def _get_children(graph_def, state_mapping, index_ref, index_ref_cache):
                             variable.update_from_ref(value)
                         elif isinstance(value, State):
                             if value._been_writen:
-                                variable.write_value(value.value)
+                                variable.value = value.value
                             else:
                                 variable.restore_value(value.value)
                         else:

brainstate/nn/_dyn_impl/_inputs.py

@@ -216,9 +216,37 @@ 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`.
+    Generates Poisson-distributed input spikes to a target state variable.
+
+    This function simulates Poisson input to a given state, updating the target
+    variable with generated spikes based on the specified frequency, number of inputs,
+    and synaptic weight. The input can be applied to specific indices of the target
+    or to the entire target if indices are not provided.
+
+    Parameters
+    ----------
+    freq : u.Quantity[u.Hz]
+        The frequency of the Poisson input in Hertz.
+    num_input : int
+        The number of input channels or neurons generating the Poisson spikes.
+    weight : u.Quantity
+        The synaptic weight applied to each spike.
+    target : State
+        The target state variable to which the Poisson input is applied.
+    indices : Optional[Union[np.ndarray, jax.Array]], optional
+        Specific indices of the target to apply the input. If None, the input is applied
+        to the entire target.
+    refractory : Optional[Union[jax.Array]], optional
+        A boolean array indicating which parts of the target are in a refractory state
+        and should not be updated. Should be the same length as the target.
+
+    Returns
+    -------
+    None
+        The function updates the target state in place with the generated Poisson input.
     """
     freq = maybe_state(freq)
     weight = maybe_state(weight)
@@ -291,7 +319,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 +386,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/nn/_interaction/_conv.py

@@ -193,16 +193,18 @@ class _Conv(_BaseConv):
         name: str = None,
         param_type: type = ParamState,
     ):
-        super().__init__(in_size=in_size,
-                         out_channels=out_channels,
-                         kernel_size=kernel_size,
-                         stride=stride,
-                         padding=padding,
-                         lhs_dilation=lhs_dilation,
-                         rhs_dilation=rhs_dilation,
-                         groups=groups,
-                         w_mask=w_mask,
-                         name=name)
+        super().__init__(
+            in_size=in_size,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            lhs_dilation=lhs_dilation,
+            rhs_dilation=rhs_dilation,
+            groups=groups,
+            w_mask=w_mask,
+            name=name
+        )
 
         self.w_initializer = w_init
         self.b_initializer = b_init

{brainstate-0.1.0.post20250209.dist-info → brainstate-0.1.0.post20250211.dist-info}/METADATA

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

@@ -1,5 +1,5 @@
 brainstate/__init__.py,sha256=AkZyyFkn4fB8g2aT6Rc2MO1xICPpUZuDtdze-eUQNc0,1496
-brainstate/_state.py,sha256=W1Q_RAL01rUSLZuOARMuX9I-26tBuIl_VzNWAziz6A8,27518
+brainstate/_state.py,sha256=TrlgaX_Hu3aSdkOUqw88HZMKUKM_PgnE9YGp595nfi0,27761
 brainstate/_state_test.py,sha256=rJUFRSXEqrrl4qANRewY9mnDlzSbtHwBIGeZ0ku-8Dg,1650
 brainstate/_utils.py,sha256=uJ6WWKq3yb05ZdktCQGLWOXsOJveL1H9pR7eev70Jes,1693
 brainstate/environ.py,sha256=PZnVFWPioUBuWmwCO8wwCKrHQfP3BR-5lYPRl5i5GDA,17698
@@ -24,17 +24,17 @@ brainstate/compile/_conditions.py,sha256=gApsHKGQrf1QBjoKXDVL7VsoeJ2zFtSc-hFz9nb
 brainstate/compile/_conditions_test.py,sha256=s9LF6h9LvigvgxUIugTqvgCHBIU8TXS1Ar1OlIxXfrw,8389
 brainstate/compile/_error_if.py,sha256=TFvhqITKkRO9m30GdlUP4eEjJvLWQUhjkujXO9zvrWs,2689
 brainstate/compile/_error_if_test.py,sha256=OdJG483IIdOrCHxtHd49OHfOxCSnSkk7GdAUOzSt8bE,2044
-brainstate/compile/_jit.py,sha256=itAWENKfJvnlaWl_uSy8lHTK8K1in89F_ZXXwp-EGRM,13944
+brainstate/compile/_jit.py,sha256=-Y8fyy8gc7qQT2ti4-N-74hjP_6C-D8YC5h-1unEKuI,13910
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
 brainstate/compile/_loop_collect_return.py,sha256=TrKBZhtQecTtuiVz_HOeyepde-znzjlyk0Te53-AvOE,23492
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
 brainstate/compile/_loop_no_collection.py,sha256=qto2__Zt2PJntkjB9AXEgraGLvNUJS483BhCXjJyqv0,7495
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
-brainstate/compile/_make_jaxpr.py,sha256=MuAa9LjXi29DjYgDUrK0WaomkjbhHZk9mWW04XGcV98,33297
+brainstate/compile/_make_jaxpr.py,sha256=Rr36U0s8qow1A4KJYXkALX10Rm2pkSYF2j_1eiSuSGI,33292
 brainstate/compile/_make_jaxpr_test.py,sha256=3gwdiutn_PJyiweu3oPEXumxEVHKaE2xDGvkwZy2GEo,4367
 brainstate/compile/_progress_bar.py,sha256=5pCMCEmbTO5XmKtzRUJGA178tuBznWKuh9Kw00wAL1I,7524
 brainstate/compile/_unvmap.py,sha256=CJA6D9lUcBfvdLrpFVvC2AdTJqe9uY0Ht6PltQJyr4U,4228
-brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
+brainstate/compile/_util.py,sha256=iKk51BHAerBFj2BTxPNdjsk3MZQiXenzpCr7Ys0iYWg,6299
 brainstate/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
 brainstate/functional/_activations.py,sha256=S0Ok7sq5FTbmJWSejpOCHo1jpKX0gYOLy_TO2IUXM8s,21726
 brainstate/functional/_activations_test.py,sha256=T___RlSrIfXwlkw8dg5A9EZMTZGDzv3a2evUwq_nYFg,13034
@@ -44,7 +44,7 @@ brainstate/functional/_spikes.py,sha256=QY-2ayJkgkGELcq-bftPEaf_hJptVf_SP3fY36Qv
 brainstate/graph/__init__.py,sha256=noo4TjBg6iEhjjwk0sAGUhR7Ge-z8Vnc2rLYUvnqttw,1295
 brainstate/graph/_graph_node.py,sha256=XwzOuaZG9x4eZknQjzJoTnnYAy7wcKD5Vox1VkYr8GM,8345
 brainstate/graph/_graph_node_test.py,sha256=BFGfdzZFDHI0XK7hHotSVWKt3em1taGvn8FHF9NCXx8,2702
-brainstate/graph/_graph_operation.py,sha256=cIwGo3ICgtce2fmdn917r81evMFjJIKeW9doaQK4DD8,64111
+brainstate/graph/_graph_operation.py,sha256=UtBNP7hvxa-5i99LQJStXbFhUbl3icdfTq1oF4MeH1g,64106
 brainstate/graph/_graph_operation_test.py,sha256=zjvpKjQAFWtw8YZuqOk_jmlZNb_-E8oPyNx57dyc8jI,18556
 brainstate/init/__init__.py,sha256=R1dHgub47o-WJM9QkFLc7x_Q7GsyaKKDtrRHTFPpC5g,1097
 brainstate/init/_base.py,sha256=B_NLS9aKNrvuj5NAlSgBbQTVev7IRvzcx8vH0J-Gq2w,1671
@@ -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=x4bcp7fo5SI5TC4TmyARngK-PE7OvGHprJ17Levs0ls,12937
 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
@@ -84,7 +84,7 @@ brainstate/nn/_elementwise/_dropout_test.py,sha256=k6aB5v8RYMoV5w8UV9UNSFhaQTV7w
 brainstate/nn/_elementwise/_elementwise.py,sha256=om-KpwDTk5yFG5KBYXXHquRLV7s28_FJjk-omvyMyvQ,33342
 brainstate/nn/_elementwise/_elementwise_test.py,sha256=SZI9jB39sZ5SO1dpWGW-PhodthwN0GU9FY1nqf2fWcs,5341
 brainstate/nn/_interaction/__init__.py,sha256=TTY_SeNrdx4VnUSw6vdyl02OHdS9Qs15cWBp6kjsyNQ,1289
-brainstate/nn/_interaction/_conv.py,sha256=lwyxTVsJVPiKlZcgB6iqE64aX7AOJzplDSj4y6-m18o,18592
+brainstate/nn/_interaction/_conv.py,sha256=eKhABWtG3QlOy7TPY9yoQjP3liBh9bb8X5Wns3_YUUQ,18499
 brainstate/nn/_interaction/_conv_test.py,sha256=fHXRFYnDghFiKre63RqMwIE_gbPKdK34UPhKOz-J3qU,8695
 brainstate/nn/_interaction/_embedding.py,sha256=iK0I1ExKWFa_QzV9UDGj32Ljsmdr1g_LlAtMcusebxU,2187
 brainstate/nn/_interaction/_linear.py,sha256=EnkOk1oE79rvRIjU6HBllxUpVOEcQQCj4vtavo9AJjI,14767
@@ -117,8 +117,8 @@ brainstate/util/_others.py,sha256=jsPZwP-v_5HRV-LB5F0NUsiqr04y8bmGIsu_JMyVcbQ,14
 brainstate/util/_pretty_repr.py,sha256=vNwRlj4sI4QJ_koyIs7eKdUMeB_QWwzRYsE8PpAWN3g,5833
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=KMMHcshOM20gYhSahNzWLxsTt-Rt3AeX3Uz26-rP9vI,17619
-brainstate-0.1.0.post20250209.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250209.dist-info/METADATA,sha256=vc9kKmrq5JM9Os6brL4zecy55nEpd9ASK9GZNJBQV9g,3585
-brainstate-0.1.0.post20250209.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250209.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250209.dist-info/RECORD,,
+brainstate-0.1.0.post20250211.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250211.dist-info/METADATA,sha256=werv_oEsECW5xLgvO__Yjth5vLfz-YARh442Q6E6FIk,3585
+brainstate-0.1.0.post20250211.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250211.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250211.dist-info/RECORD,,