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

brainstate/compile/_jit.py

@@ -40,7 +40,7 @@ class JittedFunction(Callable):
     jitted_fun: jax.stages.Wrapped  # the jitted function
     clear_cache: Callable  # clear the cache of the jitted function
     eval_shape: Callable  # evaluate the shape of the jitted function
-    lower: Callable  # lower the jitted function
+    compile: Callable  # lower the jitted function
     trace: Callable  # trace the jitted
 
     def __call__(self, *args, **kwargs):
@@ -104,7 +104,18 @@ def _get_jitted_fun(
     def eval_shape():
         raise NotImplementedError
 
-    def lower():
+    def trace():
+        """Trace this function explicitly for the given arguments.
+
+        A traced function is staged out of Python and translated to a jaxpr. It is
+        ready for lowering but not yet lowered.
+
+        Returns:
+          A ``Traced`` instance representing the tracing.
+        """
+        raise NotImplementedError
+
+    def compile(*args, **params):
         """Lower this function explicitly for the given arguments.
 
         A lowered function is staged out of Python and translated to a
@@ -114,18 +125,13 @@ def _get_jitted_fun(
         Returns:
           A ``Lowered`` instance representing the lowering.
         """
-        raise NotImplementedError
-
-    def trace():
-        """Trace this function explicitly for the given arguments.
+        # compile the function and get the state trace
+        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)
 
-        A traced function is staged out of Python and translated to a jaxpr. It is
-        ready for lowering but not yet lowered.
+        # call the jitted function
+        return jit_fun.lower(state_trace.get_state_values(), *args, **params).compile()
 
-        Returns:
-          A ``Traced`` instance representing the tracing.
-        """
-        raise NotImplementedError
 
     jitted_fun: JittedFunction
 
@@ -144,8 +150,8 @@ def _get_jitted_fun(
     # evaluate the shape of the jitted function
     jitted_fun.eval_shape = eval_shape
 
-    # lower the jitted function
-    jitted_fun.lower = lower
+    # compile the jitted function
+    jitted_fun.compile = compile
 
     # trace the jitted
     jitted_fun.trace = trace

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',
+]