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

brainstate/compile/_make_jaxpr.py

@@ -730,7 +730,7 @@ def _make_jaxpr(
     def _abstractify(args, kwargs):
         flat_args, in_tree = jax.tree.flatten((args, kwargs))
         if abstracted_axes is None:
-            return map(jax.api_util.shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
+            return map(shaped_abstractify, flat_args), in_tree, [True] * len(flat_args)
         else:
             axes_specs = _flat_axes_specs(abstracted_axes, *args, **kwargs)
             in_type = pe.infer_lambda_input_type(axes_specs, flat_args)

brainstate/compile/_progress_bar.py

@@ -35,6 +35,47 @@ Output = Any
 class ProgressBar(object):
     """
     A progress bar for tracking the progress of a jitted for-loop computation.
+
+    It can be used in :py:func:`for_loop`, :py:func:`checkpointed_for_loop`, :py:func:`scan`,
+    and :py:func:`checkpointed_scan` functions. Or any other jitted function that uses
+    a for-loop.
+
+    The message displayed in the progress bar can be customized by the following two methods:
+
+    1. By passing a string to the `desc` argument. For example:
+
+    .. code-block:: python
+
+            ProgressBar(desc="Running 1000 iterations")
+
+    2. By passing a tuple with a string and a callable function to the `desc` argument. The callable
+       function should take a dictionary as input and return a dictionary. The returned dictionary
+       will be used to format the string. For example:
+
+    .. code-block:: python
+
+                a = bst.State(1.)
+                def loop_fn(x):
+                    a.value = x.value + 1.
+                    return jnp.sum(x ** 2)
+
+                pbar = ProgressBar(desc=("Running {i} iterations, loss = {loss}",
+                                         lambda i_carray_y: {"i": i_carray_y["i"], "loss": i_carray_y["y"]}))
+
+                bst.compile.for_loop(loop_fn, xs, pbar=pbar)
+
+    In this example, ``"i"`` denotes the iteration number and ``"loss"`` is computed from the output,
+    the ``"carry"`` is the dynamic state in the loop, for example ``a.value`` in this case.
+
+
+    Args:
+        freq: The frequency at which to print the progress bar. If not specified, the progress
+            bar will be printed every 5% of the total iterations.
+        count: The number of times to print the progress bar. If not specified, the progress
+            bar will be printed every 5% of the total iterations.
+        desc: A description of the progress bar. If not specified, a default message will be
+            displayed.
+        kwargs: Additional keyword arguments to pass to the progress bar.
     """
     __module__ = "brainstate.compile"
 
@@ -42,7 +83,7 @@ class ProgressBar(object):
         self,
         freq: Optional[int] = None,
         count: Optional[int] = None,
-        desc: Optional[Tuple[str, Callable[[Dict], Dict]]] = None,
+        desc: Optional[Tuple[str, Callable[[Dict], Dict]] | str] = None,
         **kwargs
     ):
         # print rate
@@ -62,9 +103,12 @@ class ProgressBar(object):
 
         # description
         if desc is not None:
-            assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
-            assert isinstance(desc[0], str), 'Description should be a string.'
-            assert callable(desc[1]), 'Description should be a callable.'
+            if isinstance(desc, str):
+                pass
+            else:
+                assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
+                assert isinstance(desc[0], str), 'Description should be a string.'
+                assert callable(desc[1]), 'Description should be a callable.'
         self.desc = desc
 
         # check if tqdm is installed
@@ -136,8 +180,7 @@ class ProgressBarRunner(object):
         self.tqdm_bars[0].close()
 
     def __call__(self, iter_num, **kwargs):
-        data = dict(i=iter_num, **kwargs)
-        data = dict() if isinstance(self.message, str) else self.message[1](data)
+        data = dict() if isinstance(self.message, str) else self.message[1](dict(i=iter_num, **kwargs))
         assert isinstance(data, dict), 'Description function should return a dictionary.'
 
         _ = jax.lax.cond(

brainstate/nn/_dyn_impl/_inputs.py

@@ -216,15 +216,21 @@ 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)
 
     assert isinstance(target, State), 'The target must be a State.'
-    p = (freq * environ.get_dt()).to_decimal()
+    p = freq * environ.get_dt()
+    p = p.to_decimal() if isinstance(p, u.Quantity) else p
     a = num_input * p
     b = num_input * (1 - p)
     tar_val = target.value
@@ -290,8 +296,9 @@ def poisson_input(
         # )
 
         # update target variable
-        target.value = jax.tree.map(
-            lambda x: x * weight,
+        data = jax.tree.map(
+            lambda tar, x: tar + x * weight,
+            target.value,
             inp,
             is_leaf=u.math.is_quantity
         )
@@ -356,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.post20250126.dist-info → brainstate-0.1.0.post20250129.dist-info}/METADATA

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

@@ -30,9 +30,9 @@ brainstate/compile/_loop_collect_return.py,sha256=TrKBZhtQecTtuiVz_HOeyepde-znzj
 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=BwQZXfpWeXim2Pq6I7NgIvZ8p0hSH_QFyVlR4CHswmw,33322
+brainstate/compile/_make_jaxpr.py,sha256=J4oWoPBwG-fdJvNhBEtNgmo3rXrIWCoajELhaIumgPU,33309
 brainstate/compile/_make_jaxpr_test.py,sha256=3gwdiutn_PJyiweu3oPEXumxEVHKaE2xDGvkwZy2GEo,4367
-brainstate/compile/_progress_bar.py,sha256=0oVlZ4kW_ZMciJjOR_ebj3PNe_XkCMkoQpv-HUUdoF0,5554
+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/functional/__init__.py,sha256=j6-3Er4fgqWpvntzYCZVB3e5hoz-Z3aqvapITCuDri0,1107
@@ -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=UNoGxKIKXwPnhelljDowqAWlV6ds7aBBkEbvdy2oDI4,11302
+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.post20250126.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250126.dist-info/METADATA,sha256=G9QbLQwKk1SwCQ4P2MhFRufKVefEUHsl_DNDut8GRdQ,3585
-brainstate-0.1.0.post20250126.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250126.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250126.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,,