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

brainstate/compile/_conditions.py

@@ -94,9 +94,8 @@ def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
             return false_fun(*operands)
 
     # evaluate jaxpr
-    with jax.ensure_compile_time_eval():
-        stateful_true = StatefulFunction(true_fun).make_jaxpr(*operands)
-        stateful_false = StatefulFunction(false_fun).make_jaxpr(*operands)
+    stateful_true = StatefulFunction(true_fun).make_jaxpr(*operands)
+    stateful_false = StatefulFunction(false_fun).make_jaxpr(*operands)
 
     # state trace and state values
     state_trace = stateful_true.get_state_trace() + stateful_false.get_state_trace()
@@ -175,10 +174,9 @@ def switch(index, branches: Sequence[Callable], *operands):
         return branches[int(index)](*operands)
 
     # evaluate jaxpr
-    with jax.ensure_compile_time_eval():
-        wrapped_branches = [StatefulFunction(branch) for branch in branches]
-        for wrapped_branch in wrapped_branches:
-            wrapped_branch.make_jaxpr(*operands)
+    wrapped_branches = [StatefulFunction(branch) for branch in branches]
+    for wrapped_branch in wrapped_branches:
+        wrapped_branch.make_jaxpr(*operands)
 
     # wrap the functions
     state_trace = wrapped_branches[0].get_state_trace() + wrapped_branches[1].get_state_trace()

brainstate/compile/_jit.py

@@ -83,9 +83,9 @@ def _get_jitted_fun(
             return fun.fun(*args, **params)
 
         # compile the function and get the state trace
-        with jax.ensure_compile_time_eval():
-            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)
+        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)
+
         # call the jitted function
         write_state_vals, outs = jit_fun(state_trace.get_state_values(), *args, **params)
         # write the state values back to the states

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
@@ -202,12 +208,11 @@ def scan(
     # ------------------------------ #
     xs_avals = [jax.core.raise_to_shaped(jax.core.get_aval(x)) for x in xs_flat]
     x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
-    with jax.ensure_compile_time_eval():
-        stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
-        state_trace = stateful_fun.get_state_trace()
-        all_writen_state_vals = state_trace.get_write_state_values(True)
-        all_read_state_vals = state_trace.get_read_state_values(True)
-        wrapped_f = wrap_single_fun(stateful_fun, state_trace.been_writen, all_read_state_vals)
+    stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
+    state_trace = stateful_fun.get_state_trace()
+    all_writen_state_vals = state_trace.get_write_state_values(True)
+    all_read_state_vals = state_trace.get_read_state_values(True)
+    wrapped_f = wrap_single_fun(stateful_fun, state_trace.been_writen, all_read_state_vals)
 
     # scan
     init = (all_writen_state_vals, init)
@@ -231,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.
@@ -289,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
 
@@ -381,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`.
@@ -433,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/_loop_no_collection.py

@@ -103,11 +103,10 @@ def while_loop(
             pass
 
     # evaluate jaxpr
-    with jax.ensure_compile_time_eval():
-        stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
-        stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
-        if len(stateful_cond.get_write_states()) != 0:
-            raise ValueError("while_loop: cond_fun should not have any write states.")
+    stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
+    stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
+    if len(stateful_cond.get_write_states()) != 0:
+        raise ValueError("while_loop: cond_fun should not have any write states.")
 
     # state trace and state values
     state_trace = stateful_cond.get_state_trace() + stateful_body.get_state_trace()

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.")
@@ -105,25 +107,26 @@ class ProgressBarRunner(object):
             self.tqdm_bars[0].close()
 
     def __call__(self, iter_num, *args, **kwargs):
-        jax.debug.callback(
-            self._tqdm,
-            iter_num == 0,
-            (iter_num + 1) % self.print_freq == 0,
-            iter_num == self.n - 1
-        )
-
-        # _ = jax.lax.cond(
+        # jax.debug.callback(
+        #     self._tqdm,
         #     iter_num == 0,
-        #     lambda: jax.debug.callback(self._define_tqdm, ordered=True),
-        #     lambda: None,
-        # )
-        # _ = jax.lax.cond(
         #     (iter_num + 1) % self.print_freq == 0,
-        #     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, ordered=True),
-        #     lambda: None,
+        #     iter_num == self.n - 1
         # )
+
+        _ = jax.lax.cond(
+            iter_num == 0,
+            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, ordered=True),
+            lambda: None,
+        )
+        _ = jax.lax.cond(
+            iter_num == self.n - 1,
+            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',
+]