brainstate 0.1.0.post20241210__py2.py3-none-any.whl → 0.1.0.post20241219__py2.py3-none-any.whl

brainstate/compile/_loop_collect_return.py

@@ -63,7 +63,7 @@ def scan(
     length: int | None = None,
     reverse: bool = False,
     unroll: int | bool = 1,
-    pbar: ProgressBar | None = None,
+    pbar: ProgressBar | int | None = None,
 ) -> Tuple[Carry, Y]:
     """
     Scan a function over leading array axes while carrying along state.
@@ -177,7 +177,13 @@ def scan(
     has_pbar = False
     if pbar is not None:
         has_pbar = True
-        f = _wrap_fun_with_pbar(f, pbar.init(length))
+        if isinstance(pbar, ProgressBar):
+            pbar_runner = pbar.init(length)
+        elif isinstance(pbar, int):
+            pbar_runner = ProgressBar(freq=pbar).init(length)
+        else:
+            raise TypeError("pbar argument should be a ProgressBar instance or an integer.")
+        f = _wrap_fun_with_pbar(f, pbar_runner)
         init = (0, init) if pbar else init
 
     # not jit
@@ -230,7 +236,7 @@ def checkpointed_scan(
     xs: X,
     length: Optional[int] = None,
     base: int = 16,
-    pbar: Optional[ProgressBar] = None,
+    pbar: Optional[ProgressBar | int] = None,
 ) -> Tuple[Carry, Y]:
     """
     Scan a function over leading array axes while carrying along state.
@@ -288,8 +294,10 @@ def checkpointed_scan(
             length, = unique_lengths
 
     # function with progress bar
-    if pbar is not None:
+    if isinstance(pbar, ProgressBar):
         pbar_runner = pbar.init(length)
+    elif isinstance(pbar, int):
+        pbar_runner = ProgressBar(freq=pbar).init(length)
     else:
         pbar_runner = None
 
@@ -380,7 +388,7 @@ def for_loop(
     length: Optional[int] = None,
     reverse: bool = False,
     unroll: int | bool = 1,
-    pbar: Optional[ProgressBar] = None
+    pbar: Optional[ProgressBar | int] = None
 ) -> Y:
     """
     ``for-loop`` control flow with :py:class:`~.State`.
@@ -432,7 +440,7 @@ def checkpointed_for_loop(
     *xs: X,
     length: Optional[int] = None,
     base: int = 16,
-    pbar: Optional[ProgressBar] = None,
+    pbar: Optional[ProgressBar | int] = None,
 ) -> Y:
     """
     ``for-loop`` control flow with :py:class:`~.State` with a checkpointed version, similar to :py:func:`for_loop`.

brainstate/compile/_progress_bar.py

@@ -35,6 +35,8 @@ class ProgressBar(object):
 
     def __init__(self, freq: Optional[int] = None, count: Optional[int] = None, **kwargs):
         self.print_freq = freq
+        if isinstance(freq, int):
+            assert freq > 0, "Print rate should be > 0."
         self.print_count = count
         if self.print_freq is not None and self.print_count is not None:
             raise ValueError("Cannot specify both count and freq.")
@@ -114,17 +116,17 @@ class ProgressBarRunner(object):
 
         _ = jax.lax.cond(
             iter_num == 0,
-            lambda: jax.debug.callback(self._define_tqdm),
+            lambda: jax.debug.callback(self._define_tqdm, ordered=True),
             lambda: None,
         )
         _ = jax.lax.cond(
             iter_num % self.print_freq == (self.print_freq - 1),
-            lambda: jax.debug.callback(self._update_tqdm),
+            lambda: jax.debug.callback(self._update_tqdm, ordered=True),
             lambda: None,
         )
         _ = jax.lax.cond(
             iter_num == self.n - 1,
-            lambda: jax.debug.callback(self._close_tqdm),
+            lambda: jax.debug.callback(self._close_tqdm, ordered=True),
             lambda: None,
         )
 

brainstate/event/__init__.py

@@ -14,14 +14,16 @@
 # ==============================================================================
 
 
+from ._csr import *
 from ._csr_mv import *
-from ._csr_mv import __all__ as __all_csr
 from ._fixedprob_mv import *
-from ._fixedprob_mv import __all__ as __all_fixed_probability
 from ._linear_mv import *
 from ._xla_custom_op import *
-from ._xla_custom_op import __all__ as __all_xla_custom_op
-from ._linear_mv import __all__ as __all_linear
 
-__all__ = __all_fixed_probability + __all_linear + __all_csr + __all_xla_custom_op
-del __all_fixed_probability, __all_linear, __all_csr, __all_xla_custom_op
+__all__ = [
+    'CSRLinear',
+    'FixedProb',
+    'XLACustomOp',
+    'CSR',
+    'CSC',
+]