brainstate 0.1.0.post20250105__py2.py3-none-any.whl → 0.1.0.post20250126__py2.py3-none-any.whl

brainstate/compile/_make_jaxpr_test.py

@@ -18,12 +18,16 @@ from __future__ import annotations
 import unittest
 
 import jax
-import jax.extend as je
 import jax.numpy as jnp
 import pytest
 
 import brainstate as bst
 
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import jaxpr_as_fun
+else:
+    from jax.extend.core import jaxpr_as_fun
+
 
 class TestMakeJaxpr(unittest.TestCase):
     def test_compar_jax_make_jaxpr(self):
@@ -85,7 +89,7 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         jaxpr, _ = bst.compile.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
-        self.assertTrue(jnp.allclose(je.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
+        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):
@@ -103,10 +107,10 @@ class TestMakeJaxpr(unittest.TestCase):
         print()
         print(jaxpr)
         print(states)
-        print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
+        print(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value))
         jaxpr = jax.make_jaxpr(ffa)(jnp.zeros(1))
         print(jaxpr)
-        print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
     def test_compar_jax_make_jaxpr3(self):
         def fa(x):
@@ -116,10 +120,10 @@ class TestMakeJaxpr(unittest.TestCase):
         print()
         print(jaxpr)
         print(states)
-        # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
         jaxpr = jax.make_jaxpr(fa)(jnp.zeros(1))
         print(jaxpr)
-        # print(jax.core.jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
+        # print(jaxpr_as_fun(jaxpr)(jnp.zeros(1)))
 
 
 def test_return_states():

brainstate/compile/_progress_bar.py

@@ -16,34 +16,59 @@
 from __future__ import annotations
 
 import copy
-from typing import Optional
+import importlib.util
+from typing import Optional, Callable, Any, Tuple, Dict
 
 import jax
 
-try:
-    from tqdm.auto import tqdm
-except (ImportError, ModuleNotFoundError):
-    tqdm = None
+tqdm_installed = importlib.util.find_spec('tqdm') is not None
 
 __all__ = [
     'ProgressBar',
 ]
 
+Index = int
+Carray = Any
+Output = Any
+
 
 class ProgressBar(object):
+    """
+    A progress bar for tracking the progress of a jitted for-loop computation.
+    """
     __module__ = "brainstate.compile"
 
-    def __init__(self, freq: Optional[int] = None, count: Optional[int] = None, **kwargs):
+    def __init__(
+        self,
+        freq: Optional[int] = None,
+        count: Optional[int] = None,
+        desc: Optional[Tuple[str, Callable[[Dict], Dict]]] = None,
+        **kwargs
+    ):
+        # print rate
         self.print_freq = freq
         if isinstance(freq, int):
             assert freq > 0, "Print rate should be > 0."
+
+        # print count
         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.")
+
+        # other parameters
         for kwarg in ("total", "mininterval", "maxinterval", "miniters"):
             kwargs.pop(kwarg, None)
         self.kwargs = kwargs
-        if tqdm is None:
+
+        # description
+        if desc is not None:
+            assert isinstance(desc, (tuple, list)), 'Description should be a tuple or list.'
+            assert isinstance(desc[0], str), 'Description should be a string.'
+            assert callable(desc[1]), 'Description should be a callable.'
+        self.desc = desc
+
+        # check if tqdm is installed
+        if not tqdm_installed:
             raise ImportError("tqdm is not installed.")
 
     def init(self, n: int):
@@ -67,15 +92,22 @@ class ProgressBar(object):
                 raise ValueError("Print rate should be less than the "
                                  f"number of steps {n}, got {freq}")
             remainder = n % freq
-        desc = kwargs.pop("desc", f"Running for {n:,} iterations")
-        message = kwargs.pop("message", desc)
-        return ProgressBarRunner(n, message, freq, remainder, **kwargs)
+
+        message = f"Running for {n:,} iterations" if self.desc is None else self.desc
+        return ProgressBarRunner(n, freq, remainder, message, **kwargs)
 
 
 class ProgressBarRunner(object):
     __module__ = "brainstate.compile"
 
-    def __init__(self, n: int, message, print_freq: int, remainder: int, **kwargs):
+    def __init__(
+        self,
+        n: int,
+        print_freq: int,
+        remainder: int,
+        message: str | Tuple[str, Callable[[Dict], Dict]],
+        **kwargs
+    ):
         self.tqdm_bars = {}
         self.kwargs = kwargs
         self.n = n
@@ -83,50 +115,46 @@ class ProgressBarRunner(object):
         self.remainder = remainder
         self.message = message
 
-    def _define_tqdm(self):
+    def _define_tqdm(self, x: dict):
+        from tqdm.auto import tqdm
         self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
-        self.tqdm_bars[0].set_description(self.message, refresh=False)
+        if isinstance(self.message, str):
+            self.tqdm_bars[0].set_description(self.message, refresh=False)
+        else:
+            self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
 
-    def _update_tqdm(self):
+    def _update_tqdm(self, x: dict):
         self.tqdm_bars[0].update(self.print_freq)
+        if not isinstance(self.message, str):
+            self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
 
-    def _close_tqdm(self):
+    def _close_tqdm(self, x: dict):
         if self.remainder > 0:
             self.tqdm_bars[0].update(self.remainder)
+            if not isinstance(self.message, str):
+                self.tqdm_bars[0].set_description(self.message[0].format(**x), refresh=True)
         self.tqdm_bars[0].close()
 
-    def _tqdm(self, is_init, is_print, is_final):
-        if is_init:
-            self.tqdm_bars[0] = tqdm(range(self.n), **self.kwargs)
-            self.tqdm_bars[0].set_description(self.message, refresh=False)
-        if is_print:
-            self.tqdm_bars[0].update(self.print_freq)
-        if is_final:
-            if self.remainder > 0:
-                self.tqdm_bars[0].update(self.remainder)
-            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
-        # )
+    def __call__(self, iter_num, **kwargs):
+        data = dict(i=iter_num, **kwargs)
+        data = dict() if isinstance(self.message, str) else self.message[1](data)
+        assert isinstance(data, dict), 'Description function should return a dictionary.'
 
         _ = jax.lax.cond(
             iter_num == 0,
-            lambda: jax.debug.callback(self._define_tqdm, ordered=True),
-            lambda: None,
+            lambda x: jax.debug.callback(self._define_tqdm, x, ordered=True),
+            lambda x: None,
+            data
         )
         _ = jax.lax.cond(
             iter_num % self.print_freq == (self.print_freq - 1),
-            lambda: jax.debug.callback(self._update_tqdm, ordered=True),
-            lambda: None,
+            lambda x: jax.debug.callback(self._update_tqdm, x, ordered=True),
+            lambda x: None,
+            data
         )
         _ = jax.lax.cond(
             iter_num == self.n - 1,
-            lambda: jax.debug.callback(self._close_tqdm, ordered=True),
-            lambda: None,
+            lambda x: jax.debug.callback(self._close_tqdm, x, ordered=True),
+            lambda x: None,
+            data
         )
-

brainstate/compile/_unvmap.py

@@ -16,13 +16,16 @@ from __future__ import annotations
 
 import jax
 import jax.core
-import jax.extend as je
 import jax.interpreters.batching as batching
 import jax.interpreters.mlir as mlir
 import jax.numpy as jnp
-
 from brainstate._utils import set_module_as
 
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import Primitive
+else:
+    from jax.extend.core import Primitive
+
 __all__ = [
     "unvmap",
 ]
@@ -44,7 +47,7 @@ def unvmap(x, op: str = 'any'):
 
 # unvmap_all
 
-unvmap_all_p = je.core.Primitive("unvmap_all")
+unvmap_all_p = Primitive("unvmap_all")
 
 
 def unvmap_all(x):
@@ -75,7 +78,7 @@ mlir.register_lowering(
 
 # unvmap_any
 
-unvmap_any_p = jax.core.Primitive("unvmap_any")
+unvmap_any_p = Primitive("unvmap_any")
 
 
 def unvmap_any(x):
@@ -106,7 +109,7 @@ mlir.register_lowering(
 
 # unvmap_max
 
-unvmap_max_p = jax.core.Primitive("unvmap_max")
+unvmap_max_p = Primitive("unvmap_max")
 
 
 def unvmap_max(x):
@@ -153,7 +156,7 @@ def _without_vmap_batch(x, batch_axes):
     return _without_vmap(x), batching.not_mapped
 
 
-_no_vmap_prim = jax.core.Primitive('no_vmap')
+_no_vmap_prim = Primitive('no_vmap')
 _no_vmap_prim.def_impl(_without_vmap_imp)
 _no_vmap_prim.def_abstract_eval(_without_vmap_abs)
 batching.primitive_batchers[_no_vmap_prim] = _without_vmap_batch

brainstate/graph/__init__.py

@@ -14,20 +14,16 @@
 # ==============================================================================
 
 
-from ._graph_context import *
-from ._graph_context import __all__ as _graph_context__all__
-from ._graph_convert import *
-from ._graph_convert import __all__ as _graph_convert__all__
-from ._graph_node import *
-from ._graph_node import __all__ as _graph_node__all__
-from ._graph_operation import *
-from ._graph_operation import __all__ as _graph_operation__all__
+from ._graph_node import Node, Dict, List, Sequential
+from ._graph_operation import (
+    pop_states, nodes, states, treefy_states, update_states, flatten, unflatten,
+    treefy_split, treefy_merge, iter_leaf, iter_node, clone, graphdef,
+    call, RefMap, GraphDef, NodeRef, NodeDef
+)
 
-__all__ = (_graph_context__all__ +
-           _graph_convert__all__ +
-           _graph_node__all__ +
-           _graph_operation__all__)
-del (_graph_context__all__,
-     _graph_convert__all__,
-     _graph_node__all__,
-     _graph_operation__all__)
+__all__ = [
+    'Node', 'Dict', 'List', 'Sequential',
+    'pop_states', 'nodes', 'states', 'treefy_states', 'update_states', 'flatten', 'unflatten',
+    'treefy_split', 'treefy_merge', 'iter_leaf', 'iter_node', 'clone', 'graphdef',
+    'call', 'RefMap', 'GraphDef', 'NodeRef', 'NodeDef',
+]

brainstate/graph/_graph_node.py

@@ -1,7 +1,7 @@
 # The file is adapted from the Flax library (https://github.com/google/flax).
 # The credit should go to the Flax authors.
 #
-# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,9 +27,7 @@ import numpy as np
 
 from brainstate._state import State, TreefyState
 from brainstate.typing import Key
-from brainstate.util._error import TraceContextError
 from brainstate.util._pretty_repr import PrettyRepr, pretty_repr_avoid_duplicate, PrettyType, PrettyAttr
-from brainstate.util._tracers import StateJaxTracer
 from ._graph_operation import register_graph_node_type
 
 __all__ = [
@@ -44,7 +42,6 @@ class GraphNodeMeta(ABCMeta):
     if not TYPE_CHECKING:
         def __call__(cls, *args: Any, **kwargs: Any) -> Any:
             node = cls.__new__(cls, *args, **kwargs)
-            vars(node)['_trace_state'] = StateJaxTracer()
             node.__init__(*args, **kwargs)
             return node
 
@@ -64,9 +61,6 @@ class Node(PrettyRepr, metaclass=GraphNodeMeta):
 
     graph_invisible_attrs = ()
 
-    if TYPE_CHECKING:
-        _trace_state: StateJaxTracer
-
     def __init_subclass__(cls) -> None:
         super().__init_subclass__()
 
@@ -79,21 +73,6 @@ class Node(PrettyRepr, metaclass=GraphNodeMeta):
             clear=_node_clear,
         )
 
-    # if not TYPE_CHECKING:
-    #   def __setattr__(self, name: str, value: Any) -> None:
-    #     self._setattr(name, value)
-
-    # def _setattr(self, name: str, value: Any) -> None:
-    #   self.check_valid_context(lambda: f"Cannot mutate '{type(self).__name__}' from different trace level")
-    #   object.__setattr__(self, name, value)
-
-    def check_valid_context(self, error_msg: Callable[[], str]) -> None:
-        """
-        Check if the current context is valid for the object to be mutated.
-        """
-        if not self._trace_state.is_valid():
-            raise TraceContextError(error_msg())
-
     def __deepcopy__(self: G, memo=None) -> G:
         """
         Deepcopy the object.
@@ -214,7 +193,6 @@ def _node_create_empty(
 ) -> G:
     node_type, = static
     node = object.__new__(node_type)
-    vars(node).update(_trace_state=StateJaxTracer())
     return node
 
 

brainstate/graph/_graph_operation.py

@@ -1,7 +1,7 @@
 # The file is adapted from the Flax library (https://github.com/google/flax).
 # The credit should go to the Flax authors.
 #
-# Copyright 2024 The Flax Authors & 2024 BDP Ecosystem.
+# Copyright 2024 The Flax Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

brainstate/graph/_graph_operation_test.py

@@ -17,13 +17,10 @@ from __future__ import annotations
 
 import unittest
 from collections.abc import Callable
-from functools import partial
 from threading import Thread
-from typing import Any
 
 import jax
 import jax.numpy as jnp
-import pytest
 from absl.testing import absltest, parameterized
 
 import brainstate as bst
@@ -354,125 +351,6 @@ class TestGraphUtils(absltest.TestCase):
         assert m2.tree.a is not m.tree.a
         assert m2.tree is not m.tree
 
-    @pytest.mark.skip(reason='Not implemented')
-    def test_cached_unflatten(self):
-        class Foo(bst.graph.Node):
-            def __init__(self, ):
-                self.a = bst.nn.Linear(2, 2)
-                self.b = bst.nn.BatchNorm1d([10, 2])
-
-        def f(m: Foo):
-            m.a, m.b = m.b, m.a  # type: ignore
-
-        m = Foo()
-        a = m.a
-        b = m.b
-
-        ref_out_idx_out = bst.graph.RefMap()
-        graphdef: bst.graph.GraphDef[Foo]
-        graphdef, state = bst.graph.flatten(m, ref_index=ref_out_idx_out)
-
-        @partial(jax.jit, static_argnums=(0,))
-        def f_pure(graphdef: bst.graph.GraphDef[Foo], state):
-            idx_out_ref_in: dict[int, Any] = {}
-            m = bst.graph.unflatten(graphdef, state, index_ref=idx_out_ref_in)
-            f(m)
-            ref_in_idx_in = bst.graph.RefMap[Any, int]()
-            graphdef, state = bst.graph.flatten(m, ref_index=ref_in_idx_in)
-            idx_out_idx_in = bst.graph.compose_mapping(idx_out_ref_in, ref_in_idx_in)
-            static_out = bst.graph.Static((graphdef, idx_out_idx_in))
-            return state, static_out
-
-        static_out: bst.graph.Static
-        state, static_out = f_pure(graphdef, state)
-        idx_out_idx_in: dict[int, int]
-        graphdef, idx_out_idx_in = static_out.value
-        idx_in_ref_out = bst.graph.compose_mapping_reversed(
-            ref_out_idx_out, idx_out_idx_in
-        )
-        m2 = bst.graph.unflatten(graphdef, state, index_ref_cache=idx_in_ref_out)
-        assert m2 is m
-        assert m2.a is b
-        assert m2.b is a
-
-    @pytest.mark.skip(reason='Not implemented')
-    def test_cached_unflatten_swap_variables(self):
-        class Foo(bst.graph.Node):
-            def __init__(self):
-                self.a = bst.ParamState(1)
-                self.b = bst.ParamState(2)
-
-        def f(m: Foo):
-            m.a, m.b = m.b, m.a
-
-        m = Foo()
-        a = m.a
-        b = m.b
-
-        ref_out_idx_out = bst.graph.RefMap[Any, int]()
-        graphdef: bst.graph.GraphDef[Foo]
-        graphdef, state = bst.graph.flatten(m, ref_index=ref_out_idx_out)
-
-        @partial(jax.jit, static_argnums=(0,))
-        def f_pure(graphdef: bst.graph.GraphDef[Foo], state):
-            idx_out_ref_in: dict[int, Any] = {}
-            m = bst.graph.unflatten(graphdef, state, index_ref=idx_out_ref_in)
-            f(m)
-            ref_in_idx_in = bst.graph.RefMap[Any, int]()
-            graphdef, state = bst.graph.flatten(m, ref_index=ref_in_idx_in)
-            idx_out_idx_in = bst.graph.compose_mapping(idx_out_ref_in, ref_in_idx_in)
-            static_out = bst.graph.Static((graphdef, idx_out_idx_in))
-            return state, static_out
-
-        static_out: bst.graph.Static
-        state, static_out = f_pure(graphdef, state)
-        idx_out_idx_in: dict[int, int]
-        graphdef, idx_out_idx_in = static_out.value
-        idx_in_ref_out = bst.graph.compose_mapping_reversed(
-            ref_out_idx_out, idx_out_idx_in
-        )
-        m2 = bst.graph.unflatten(graphdef, state, index_ref_cache=idx_in_ref_out)
-        assert m2 is m
-        assert m2.a is b
-        assert m2.b is a
-
-    @pytest.mark.skip(reason='Not implemented')
-    def test_cached_unflatten_add_self_reference(self):
-        class Foo(bst.graph.Node):
-            def __init__(self):
-                self.ref = None
-
-        def f(m: Foo):
-            m.ref = m
-
-        m = Foo()
-
-        ref_out_idx_out = bst.graph.RefMap()
-        graphdef: bst.graph.GraphDef[Foo]
-        graphdef, state = bst.graph.flatten(m, ref_index=ref_out_idx_out)
-
-        @partial(jax.jit, static_argnums=(0,))
-        def f_pure(graphdef: bst.graph.GraphDef[Foo], state):
-            idx_out_ref_in: dict[int, Any] = {}
-            m = bst.graph.unflatten(graphdef, state, index_ref=idx_out_ref_in)
-            f(m)
-            ref_in_idx_in = bst.graph.RefMap[Any, int]()
-            graphdef, state = bst.graph.flatten(m, ref_index=ref_in_idx_in)
-            idx_out_idx_in = bst.graph.compose_mapping(idx_out_ref_in, ref_in_idx_in)
-            static_out = bst.graph.Static((graphdef, idx_out_idx_in))
-            return state, static_out
-
-        static_out: bst.graph.Static
-        state, static_out = f_pure(graphdef, state)
-        idx_out_idx_in: dict[int, int]
-        graphdef, idx_out_idx_in = static_out.value
-        idx_in_ref_out = bst.graph.compose_mapping_reversed(
-            ref_out_idx_out, idx_out_idx_in
-        )
-        m2 = bst.graph.unflatten(graphdef, state, index_ref_cache=idx_in_ref_out)
-        assert m2 is m
-        assert m2.ref is m2
-
     def test_call_jit_update(self):
         class Counter(bst.graph.Node):
             def __init__(self):
@@ -527,43 +405,6 @@ class TestGraphUtils(absltest.TestCase):
         self.assertEqual(nodes['a'].count.value, 0)
         self.assertEqual(nodes['b'].count.value, 1)
 
-    def test_to_tree_simple(self):
-        m = bst.nn.Linear(2, 3, )
-        impure_tree = (m, 1, {'b': m})
-
-        pure_tree = bst.graph.graph_to_tree(impure_tree)
-
-        t1 = pure_tree[0]
-        t2 = pure_tree[2]['b']
-
-        self.assertEqual(pure_tree[1], 1)
-        self.assertIsInstance(t1, bst.graph.NodeStates)
-        assert isinstance(t1, bst.graph.NodeStates)
-        self.assertIsInstance(t2, bst.graph.NodeStates)
-        assert isinstance(t2, bst.graph.NodeStates)
-        self.assertIsInstance(t1.graphdef, bst.graph.NodeDef)
-        self.assertIsInstance(t2.graphdef, bst.graph.NodeRef)
-        self.assertLen(t1.states[0].to_flat(), 1)
-        self.assertLen(t2.states[0].to_flat(), 0)
-
-        impure_tree2 = bst.graph.tree_to_graph(pure_tree)
-
-        m1_out = impure_tree2[0]
-        m2_out = impure_tree2[2]['b']
-
-        self.assertIs(m1_out, m2_out)
-        self.assertEqual(impure_tree2[1], 1)
-
-    def test_to_tree_consistent_prefix(self):
-        m = bst.nn.Linear(2, 3, )
-        impure_tree = (m, 1, {'b': m})
-        prefix = (0, None, 0)
-        pure_tree = bst.graph.graph_to_tree(impure_tree, prefix=prefix)
-
-        prefix = (0, None, 1)
-        with self.assertRaisesRegex(ValueError, 'Inconsistent aliasing detected'):
-            bst.graph.graph_to_tree(impure_tree, prefix=prefix)
-
 
 class SimpleModule(bst.nn.Module):
     pass