brainstate 0.1.0.post20250503__py2.py3-none-any.whl → 0.1.2__py2.py3-none-any.whl

brainstate/compile/_error_if_test.py

@@ -13,43 +13,40 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax
 import jax.numpy as jnp
-import jaxlib.xla_extension
 
-import brainstate as bst
+import brainstate
 
 
 class TestJitError(unittest.TestCase):
     def test1(self):
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-            bst.compile.jit_error_if(True, 'error')
+        with self.assertRaises(Exception):
+            brainstate.compile.jit_error_if(True, 'error')
 
         def err_f(x):
             raise ValueError(f'error: {x}')
 
-        bst.compile.jit_error_if(False, err_f, 1.)
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-            bst.compile.jit_error_if(True, err_f, 1.)
+        brainstate.compile.jit_error_if(False, err_f, 1.)
+        with self.assertRaises(Exception):
+            brainstate.compile.jit_error_if(True, err_f, 1.)
 
     def test_vmap(self):
         def f(x):
-            bst.compile.jit_error_if(x, 'error: {x}', x=x)
+            brainstate.compile.jit_error_if(x, 'error: {x}', x=x)
 
         jax.vmap(f)(jnp.array([False, False, False]))
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        with self.assertRaises(Exception):
             jax.vmap(f)(jnp.array([True, False, False]))
 
     def test_vmap_vmap(self):
         def f(x):
-            bst.compile.jit_error_if(x, 'error: {x}', x=x)
+            brainstate.compile.jit_error_if(x, 'error: {x}', x=x)
 
         jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
                                          [False, False, False]]))
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        with self.assertRaises(Exception):
             jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
                                              [True, False, False]]))

brainstate/compile/_jit.py

@@ -13,16 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import functools
 from collections.abc import Iterable, Sequence
 from typing import (Any, Callable, Union)
 
 import jax
 from jax._src import sharding_impls
-from jax.lib import xla_client as xc
 
+from brainstate._compatible_import import Device
 from brainstate._utils import set_module_as
 from brainstate.typing import Missing
 from ._make_jaxpr import StatefulFunction, _ensure_index_tuple
@@ -94,8 +92,8 @@ def _get_jitted_fun(
         read_state_vals = state_trace.get_read_state_values(True)
 
         # call the jitted function
-        # print('Running ...')
         write_state_vals, outs = jit_fun(state_trace.get_state_values(), *args, **params)
+
         # write the state values back to the states
         write_back_state_values(state_trace, read_state_vals, write_state_vals)
         return outs
@@ -106,8 +104,11 @@ def _get_jitted_fun(
         """
         # clear the cache of the stateful function
         fun.clear_cache()
-        # clear the cache of the jitted function
-        jit_fun.clear_cache()
+        try:
+            # clear the cache of the jitted function
+            jit_fun.clear_cache()
+        except AttributeError:
+            pass
 
     def eval_shape():
         raise NotImplementedError
@@ -165,7 +166,7 @@ def _get_jitted_fun(
     # compile the jitted function
     jitted_fun.compile = compile
 
-    # trace the jitted
+    # trace the jitted function
     jitted_fun.trace = trace
 
     return jitted_fun
@@ -180,7 +181,7 @@ def jit(
     donate_argnums: int | Sequence[int] | None = None,
     donate_argnames: str | Iterable[str] | None = None,
     keep_unused: bool = False,
-    device: xc.Device | None = None,
+    device: Device | None = None,
     backend: str | None = None,
     inline: bool = False,
     abstracted_axes: Any | None = None,

brainstate/compile/_loop_collect_return.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import math
 from functools import wraps
 from typing import Callable, Optional, TypeVar, Tuple, Any

brainstate/compile/_loop_collect_return_test.py

@@ -13,20 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax.numpy as jnp
 import numpy as np
 
-import brainstate as bst
+import brainstate
 
 
 class TestForLoop(unittest.TestCase):
     def test_for_loop(self):
-        a = bst.ShortTermState(0.)
-        b = bst.ShortTermState(0.)
+        a = brainstate.ShortTermState(0.)
+        b = brainstate.ShortTermState(0.)
 
         def f(i):
             a.value += (1 + b.value)
@@ -34,7 +32,7 @@ class TestForLoop(unittest.TestCase):
 
         n_iter = 10
         ops = np.arange(n_iter)
-        r = bst.compile.for_loop(f, ops)
+        r = brainstate.compile.for_loop(f, ops)
 
         print(a)
         print(b)
@@ -42,8 +40,8 @@ class TestForLoop(unittest.TestCase):
         self.assertTrue(jnp.allclose(r, ops + 1))
 
     def test_checkpointed_for_loop(self):
-        a = bst.ShortTermState(0.)
-        b = bst.ShortTermState(0.)
+        a = brainstate.ShortTermState(0.)
+        b = brainstate.ShortTermState(0.)
 
         def f(i):
             a.value += (1 + b.value)
@@ -51,7 +49,7 @@ class TestForLoop(unittest.TestCase):
 
         n_iter = 18
         ops = jnp.arange(n_iter)
-        r = bst.compile.checkpointed_for_loop(f, ops, base=2, pbar=bst.compile.ProgressBar())
+        r = brainstate.compile.checkpointed_for_loop(f, ops, base=2, pbar=brainstate.compile.ProgressBar())
 
         print(a)
         print(b)

brainstate/compile/_loop_no_collection.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import math
 from typing import Any, Callable, TypeVar
 

brainstate/compile/_loop_no_collection_test.py

@@ -13,17 +13,16 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 from unittest import TestCase
 
-import brainstate as bst
+import brainstate 
 
 
 class TestWhileLoop(TestCase):
     def test1(self):
-        a = bst.State(1.)
-        b = bst.State(20.)
+        a = brainstate.State(1.)
+        b = brainstate.State(20.)
 
         def cond(_):
             return a.value < b.value
@@ -31,13 +30,13 @@ class TestWhileLoop(TestCase):
         def body(_):
             a.value += 1.
 
-        bst.compile.while_loop(cond, body, None)
+        brainstate.compile.while_loop(cond, body, None)
 
         print(a.value, b.value)
 
     def test2(self):
-        a = bst.State(1.)
-        b = bst.State(20.)
+        a = brainstate.State(1.)
+        b = brainstate.State(20.)
 
         def cond(x):
             return a.value < b.value
@@ -46,6 +45,6 @@ class TestWhileLoop(TestCase):
             a.value += x
             return x
 
-        r = bst.compile.while_loop(cond, body, 1.)
+        r = brainstate.compile.while_loop(cond, body, 1.)
 
         print(a.value, b.value, r)

brainstate/compile/_make_jaxpr.py

@@ -51,8 +51,6 @@ function.
 
 """
 
-from __future__ import annotations
-
 import functools
 import inspect
 import operator
@@ -65,7 +63,7 @@ from jax._src import source_info_util
 from jax._src.linear_util import annotate
 from jax._src.traceback_util import api_boundary
 from jax.api_util import shaped_abstractify
-from jax.extend.linear_util import transformation_with_aux, wrap_init
+from jax.extend.linear_util import transformation_with_aux
 from jax.interpreters import partial_eval as pe
 
 from brainstate._compatible_import import (
@@ -75,6 +73,7 @@ from brainstate._compatible_import import (
     safe_zip,
     unzip2,
     wraps,
+    wrap_init,
 )
 from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
@@ -98,7 +97,7 @@ def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
         return tuple(safe_map(operator.index, x))
 
 
-def _new_arg_fn(frame, trace, aval):
+def _jax_v04_new_arg_fn(frame, trace, aval):
     """
     Transform a new argument to a tracer.
 
@@ -119,27 +118,41 @@ def _new_arg_fn(frame, trace, aval):
     return tracer
 
 
-def _new_jax_trace():
+def _jax_v04_new_jax_trace():
     main = jax.core.thread_local_state.trace_state.trace_stack.stack[-1]
     frame = main.jaxpr_stack[-1]
     trace = pe.DynamicJaxprTrace(main, jax.core.cur_sublevel())
     return frame, trace
 
 
+def _jax_v04_new_arg():
+    # Should be within the calling of ``jax.make_jaxpr()``
+    frame, trace = _jax_v04_new_jax_trace()
+    # Set the function to transform the new argument to a tracer
+    fn = functools.partial(_jax_v04_new_arg_fn, frame, trace)
+    return fn
+
+
+def _jax_new_version_new_arg():
+    trace = jax.core.trace_ctx.trace
+
+    def wrapper(x):
+        if jax.__version_info__ < (0, 6, 1):
+            return trace.new_arg(shaped_abstractify(x))
+        else:
+            return trace.new_arg(shaped_abstractify(x), source_info=source_info_util.current())
+
+    return wrapper
+
+
 def _init_state_trace_stack(name) -> StateTraceStack:
     state_trace: StateTraceStack = StateTraceStack(name=name)
 
     if jax.__version_info__ < (0, 4, 36):
-        # Should be within the calling of ``jax.make_jaxpr()``
-        frame, trace = _new_jax_trace()
-        # Set the function to transform the new argument to a tracer
-        state_trace.set_new_arg(functools.partial(_new_arg_fn, frame, trace))
-        return state_trace
-
+        state_trace.set_new_arg(_jax_v04_new_arg())
     else:
-        trace = jax.core.trace_ctx.trace
-        state_trace.set_new_arg(trace.new_arg)
-        return state_trace
+        state_trace.set_new_arg(_jax_new_version_new_arg())
+    return state_trace
 
 
 class StatefulFunction(PrettyObject):
@@ -745,7 +758,7 @@ def _make_jaxpr(
     @wraps(fun)
     @api_boundary
     def make_jaxpr_f(*args, **kwargs):
-        f = wrap_init(fun)
+        f = wrap_init(fun, (), {}, 'brainstate.compile.make_jaxpr')
         if static_argnums:
             dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
             f, args = jax.api_util.argnums_partial(f, dyn_argnums, args)
@@ -754,12 +767,12 @@ def _make_jaxpr(
         f, out_tree = _flatten_fun(f, in_tree)
         f = annotate(f, in_type)
         if jax.__version_info__ < (0, 5, 0):
-            debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
+            debug_info_ = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
         with ExitStack() as stack:
             if axis_env is not None:
                 stack.enter_context(extend_axis_env_nd(axis_env))
             if jax.__version_info__ < (0, 5, 0):
-                jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
+                jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info_)
             else:
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)
         closed_jaxpr = ClosedJaxpr(jaxpr, consts)

brainstate/compile/_make_jaxpr_test.py

@@ -21,7 +21,7 @@ import jax
 import jax.numpy as jnp
 import pytest
 
-import brainstate as bst
+import brainstate
 from brainstate._compatible_import import jaxpr_as_fun
 
 
@@ -29,10 +29,10 @@ class TestMakeJaxpr(unittest.TestCase):
     def test_compar_jax_make_jaxpr(self):
         def func4(arg):  # Arg is a pair
             temp = arg[0] + jnp.sin(arg[1]) * 3.
-            c = bst.random.rand_like(arg[0])
+            c = brainstate.random.rand_like(arg[0])
             return jnp.sum(temp + c)
 
-        key = bst.random.DEFAULT.value
+        key = brainstate.random.DEFAULT.value
         jaxpr = jax.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
         print(jaxpr)
         self.assertTrue(len(jaxpr.in_avals) == 2)
@@ -40,66 +40,66 @@ class TestMakeJaxpr(unittest.TestCase):
         self.assertTrue(len(jaxpr.out_avals) == 1)
         self.assertTrue(jnp.allclose(jaxpr.consts[0], key))
 
-        bst.random.seed(1)
-        print(bst.random.DEFAULT.value)
+        brainstate.random.seed(1)
+        print(brainstate.random.DEFAULT.value)
 
-        jaxpr2, states = bst.compile.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
+        jaxpr2, states = brainstate.compile.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
         print(jaxpr2)
         self.assertTrue(len(jaxpr2.in_avals) == 3)
         self.assertTrue(len(jaxpr2.out_avals) == 2)
         self.assertTrue(len(jaxpr2.consts) == 0)
-        print(bst.random.DEFAULT.value)
+        print(brainstate.random.DEFAULT.value)
 
     def test_StatefulFunction_1(self):
         def func4(arg):  # Arg is a pair
             temp = arg[0] + jnp.sin(arg[1]) * 3.
-            c = bst.random.rand_like(arg[0])
+            c = brainstate.random.rand_like(arg[0])
             return jnp.sum(temp + c)
 
-        fun = bst.compile.StatefulFunction(func4).make_jaxpr((jnp.zeros(8), jnp.ones(8)))
+        fun = brainstate.compile.StatefulFunction(func4).make_jaxpr((jnp.zeros(8), jnp.ones(8)))
         print(fun.get_states())
         print(fun.get_jaxpr())
 
     def test_StatefulFunction_2(self):
-        st1 = bst.State(jnp.ones(10))
+        st1 = brainstate.State(jnp.ones(10))
 
         def f1(x):
             st1.value = x + st1.value
 
         def f2(x):
-            jaxpr = bst.compile.make_jaxpr(f1)(x)
+            jaxpr = brainstate.compile.make_jaxpr(f1)(x)
             c = 1. + x
             return c
 
         def f3(x):
-            jaxpr = bst.compile.make_jaxpr(f1)(x)
+            jaxpr = brainstate.compile.make_jaxpr(f1)(x)
             c = 1.
             return c
 
         print()
-        jaxpr = bst.compile.make_jaxpr(f1)(jnp.zeros(1))
+        jaxpr = brainstate.compile.make_jaxpr(f1)(jnp.zeros(1))
         print(jaxpr)
         jaxpr = jax.make_jaxpr(f2)(jnp.zeros(1))
         print(jaxpr)
         jaxpr = jax.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
-        jaxpr, _ = bst.compile.make_jaxpr(f3)(jnp.zeros(1))
+        jaxpr, _ = brainstate.compile.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
         self.assertTrue(jnp.allclose(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
                                      f3(jnp.zeros(1))))
 
     def test_compar_jax_make_jaxpr2(self):
-        st1 = bst.State(jnp.ones(10))
+        st1 = brainstate.State(jnp.ones(10))
 
         def fa(x):
             st1.value = x + st1.value
 
         def ffa(x):
-            jaxpr, states = bst.compile.make_jaxpr(fa)(x)
+            jaxpr, states = brainstate.compile.make_jaxpr(fa)(x)
             c = 1. + x
             return c
 
-        jaxpr, states = bst.compile.make_jaxpr(ffa)(jnp.zeros(1))
+        jaxpr, states = brainstate.compile.make_jaxpr(ffa)(jnp.zeros(1))
         print()
         print(jaxpr)
         print(states)
@@ -112,7 +112,7 @@ class TestMakeJaxpr(unittest.TestCase):
         def fa(x):
             return 1.
 
-        jaxpr, states = bst.compile.make_jaxpr(fa)(jnp.zeros(1))
+        jaxpr, states = brainstate.compile.make_jaxpr(fa)(jnp.zeros(1))
         print()
         print(jaxpr)
         print(states)
@@ -125,9 +125,9 @@ class TestMakeJaxpr(unittest.TestCase):
 def test_return_states():
     import jax.numpy
 
-    a = bst.State(jax.numpy.ones(3))
+    a = brainstate.State(jax.numpy.ones(3))
 
-    @bst.compile.jit
+    @brainstate.compile.jit
     def f():
         return a
 

brainstate/compile/_progress_bar.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import copy
 import importlib.util

brainstate/compile/_unvmap.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from __future__ import annotations
 
 import jax
 import jax.core

brainstate/compile/_util.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from functools import wraps
 from typing import Sequence, Tuple
 

brainstate/environ.py

@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import contextlib
 import dataclasses
 import functools

brainstate/functional/_activations.py

@@ -18,8 +18,6 @@
 Shared neural network activations and other functions.
 """
 
-from __future__ import annotations
-
 from typing import Any, Union, Sequence
 
 import brainunit as u