brainstate 0.1.0.post20241209__py2.py3-none-any.whl → 0.1.0.post20241210__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

@@ -202,12 +202,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)

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

@@ -105,25 +105,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),
+            lambda: None,
+        )
+        _ = jax.lax.cond(
+            iter_num % self.print_freq == (self.print_freq - 1),
+            lambda: jax.debug.callback(self._update_tqdm),
+            lambda: None,
+        )
+        _ = jax.lax.cond(
+            iter_num == self.n - 1,
+            lambda: jax.debug.callback(self._close_tqdm),
+            lambda: None,
+        )
+

{brainstate-0.1.0.post20241209.dist-info → brainstate-0.1.0.post20241210.dist-info}/METADATA

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

@@ -20,19 +20,19 @@ brainstate/augment/_random.py,sha256=rkB4w4BkKsz9p8lTk31kVHvlVPJSvtGk8REn936KI_4
 brainstate/compile/__init__.py,sha256=qZZIYoyEl51IFkFu-Hb-bP3PAEHo94HlTDf57P2ze08,1858
 brainstate/compile/_ad_checkpoint.py,sha256=5zJ1ENeTU4FzRY_uNpr85NhKfuicMMjcIbhu6-bSM4k,9451
 brainstate/compile/_ad_checkpoint_test.py,sha256=R1I76nG4zIqb6g3M_VxWts7rUC1OHJCjtQhPkcbXodk,1746
-brainstate/compile/_conditions.py,sha256=ocz6sDc7Xzabz2GnRsQmS6GDps-WP-OXUd0EZTTlG0k,10217
+brainstate/compile/_conditions.py,sha256=gApsHKGQrf1QBjoKXDVL7VsoeJ2zFtSc-hFz9nbYcF0,10113
 brainstate/compile/_conditions_test.py,sha256=s9LF6h9LvigvgxUIugTqvgCHBIU8TXS1Ar1OlIxXfrw,8389
 brainstate/compile/_error_if.py,sha256=TFvhqITKkRO9m30GdlUP4eEjJvLWQUhjkujXO9zvrWs,2689
 brainstate/compile/_error_if_test.py,sha256=SJmAfosVoGd4vhfFtb1IvjeFVW914bfTccCg6DoLWYk,1992
-brainstate/compile/_jit.py,sha256=bfEszNttEtE6npqHBam1_DBlRa39fE6qP6lGaWw2amA,13750
+brainstate/compile/_jit.py,sha256=3mQ-RUFz35wceZyKE_MoR58OBL0RK_i6sHm4rWYzMLs,13698
 brainstate/compile/_jit_test.py,sha256=zD7kck9SQJGmUDolh9P4luKwQ21fBGje1Z4STTEXIuA,4135
-brainstate/compile/_loop_collect_return.py,sha256=8vDB2l0d4sIn0apspJzkhFhxjsL7reIptDeFRI9b1tc,23002
+brainstate/compile/_loop_collect_return.py,sha256=DybSBixeuxleKJV6n9FgVBDsUTmexzS0IdgWYRqp5cU,22940
 brainstate/compile/_loop_collect_return_test.py,sha256=bA-_11E8A_0jR5umEO3e409y7bb5QYDTgSL-SBaX7kQ,1802
-brainstate/compile/_loop_no_collection.py,sha256=2rSK20enkBMXPAbsCyb7PCICPNrgaSpl5jfumgWpxA0,7401
+brainstate/compile/_loop_no_collection.py,sha256=0i31gdQ7sI-d6pvnh08ttUUwdAtpx4uoYhGuf_CyL9s,7343
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
 brainstate/compile/_make_jaxpr.py,sha256=S5O9KUB3bsxoKcfptlV0MRfKA__Ija37WxkakIRL3z0,33010
 brainstate/compile/_make_jaxpr_test.py,sha256=qJUtkyj50JQ6f4UJbOLhvRdkbNn3NSKibFL9jESdQkA,4279
-brainstate/compile/_progress_bar.py,sha256=FafEbD9KzmhCCizfQoXXLw46asn9_uiuH1U5_DMtSXg,4529
+brainstate/compile/_progress_bar.py,sha256=H544Oh10SiF5ccrKHM9ay7ZHigYIhNhSQGEKbDxRJgg,4485
 brainstate/compile/_unvmap.py,sha256=ewbLLNXiI_dBsEBaVzSS0BEXNol22sd9gMzk606lSkM,4139
 brainstate/compile/_util.py,sha256=aCvkTV--g4NsqcodTdBAISt4EwgezCbKzNUV58n-Q_Y,6304
 brainstate/event/__init__.py,sha256=wOBkq7kDg90M8Y9FuoXRlSEuu1ZzbIhCJ1dHeLqN6_Q,1194
@@ -136,8 +136,8 @@ brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7
 brainstate/util/_struct.py,sha256=0exv0oOiSt1hmx20Y4J2-pCGtCTx13WcAlEYSBkyung,17640
 brainstate/util/_tracers.py,sha256=0r5T4nhxMzI79NtqroqitsdMT4YfpgV5RdYJLS5uJ0w,2285
 brainstate/util/_visualization.py,sha256=n4ZVz10z7VBqA0cKO6vyHwEMprWJgPeEqtITzDMai2Y,1519
-brainstate-0.1.0.post20241209.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20241209.dist-info/METADATA,sha256=gXsiYWSQqOJ0CWKINESG4sSpnDkcmVYgWJWeEFLTHoA,3401
-brainstate-0.1.0.post20241209.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20241209.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20241209.dist-info/RECORD,,
+brainstate-0.1.0.post20241210.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20241210.dist-info/METADATA,sha256=E6AATarjpwssXflLfA-OCkxFxqZxqJNxHZteO6UWMhw,3401
+brainstate-0.1.0.post20241210.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20241210.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20241210.dist-info/RECORD,,