brainstate 0.1.0.post20250104__py2.py3-none-any.whl → 0.1.0.post20250120__py2.py3-none-any.whl

brainstate/event/_xla_custom_op.py

@@ -9,12 +9,20 @@ from typing import Callable, Sequence, Tuple, Protocol
 import jax
 import numpy as np
 from jax import tree_util
-from jax.core import Primitive
 from jax.interpreters import batching, ad
 from jax.interpreters import xla, mlir
-from jax.lib import xla_client
 from jaxlib.hlo_helpers import custom_call
 
+if jax.__version_info__ < (0, 4, 35):
+    from jax.lib import xla_client
+else:
+    import jax.extend as je
+
+if jax.__version_info__ < (0, 4, 38):
+    from jax.core import Primitive
+else:
+    from jax.extend.core import Primitive
+
 numba_installed = importlib.util.find_spec('numba') is not None
 
 __all__ = [
@@ -143,7 +151,10 @@ def numba_cpu_custom_call_target(output_ptrs, input_ptrs):
     xla_c_rule = cfunc(sig)(new_f)
     target_name = f'numba_custom_call_{str(xla_c_rule.address)}'
     capsule = ctypes.pythonapi.PyCapsule_New(xla_c_rule.address, b"xla._CUSTOM_CALL_TARGET", None)
-    xla_client.register_custom_call_target(target_name, capsule, "cpu")
+    if jax.__version_info__ < (0, 4, 35):
+        xla_client.register_custom_call_target(target_name, capsule, "cpu")
+    else:
+        je.ffi.register_ffi_target(target_name, capsule, "cpu", api_version=0)
 
     # call
     return custom_call(
@@ -157,7 +168,7 @@ def numba_cpu_custom_call_target(output_ptrs, input_ptrs):
 
 
 def register_numba_mlir_cpu_translation_rule(
-    primitive: jax.core.Primitive,
+    primitive: Primitive,
     cpu_kernel: Callable,
     debug: bool = False
 ):
@@ -198,7 +209,7 @@ class XLACustomOp:
         transpose_translation: Callable = None,
     ):
         # primitive
-        self.primitive = jax.core.Primitive(name)
+        self.primitive = Primitive(name)
         self.primitive.multiple_results = True
 
         # abstract evaluation

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__
 
-__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__ = (
+    _graph_node__all__ +
+    _graph_operation__all__
+)
+del (
+    _graph_node__all__,
+    _graph_operation__all__
+)

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

brainstate/nn/_dyn_impl/_inputs.py

@@ -22,8 +22,8 @@ import numpy as np
 
 from brainstate import environ, init, random
 from brainstate._state import ShortTermState
-from brainstate._state import State
-from brainstate.compile import while_loop, cond
+from brainstate._state import State, maybe_state
+from brainstate.compile import while_loop
 from brainstate.nn._dynamics._dynamics_base import Dynamics, Prefetch
 from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size, DTypeLike
@@ -198,55 +198,97 @@ class PoissonInput(Module):
         self.weight = weight
 
     def update(self):
-        p = self.freq * environ.get_dt()
-        a = self.num_input * p
-        b = self.num_input * (1 - p)
-
-        target = self.target()
         target_state = getattr(self.target.module, self.target.item)
 
         # generate Poisson input
-        inp = cond(
-            u.math.logical_and(a > 5, b > 5),
-            lambda: random.normal(a, b * p, self.indices.shape),
-            lambda: random.binomial(self.num_input, p, self.indices.shape).astype(float)
+        poisson_input(
+            self.freq,
+            self.num_input,
+            self.weight,
+            target_state,
+            self.indices,
         )
 
-        # update target variable
-        target_state.value = target.at[self.indices].add(inp * self.weight)
-
 
 def poisson_input(
-    freq: ArrayLike,
+    freq: u.Quantity[u.Hz],
     num_input: int,
-    weight: ArrayLike,
+    weight: u.Quantity,
     target: State,
     indices: Optional[Union[np.ndarray, jax.Array]] = None,
 ):
     """
     Poisson Input to the given :py:class:`brainstate.State`.
     """
+    freq = maybe_state(freq)
+    weight = maybe_state(weight)
+
     assert isinstance(target, State), 'The target must be a State.'
-    p = freq * environ.get_dt()
+    p = (freq * environ.get_dt()).to_decimal()
     a = num_input * p
     b = num_input * (1 - p)
     tar_val = target.value
+    cond = u.math.logical_and(a > 5, b > 5)
+
     if indices is None:
         # generate Poisson input
-        inp = cond(
-            u.math.logical_and(a > 5, b > 5),
-            lambda: jax.tree.map(
-                lambda tar: random.normal(a, b * p, tar.shape),
-                tar_val,
-                is_leaf=u.math.is_quantity
+        branch1 = jax.tree.map(
+            lambda tar: random.normal(
+                a,
+                b * p,
+                tar.shape,
+                dtype=tar.dtype
             ),
-            lambda: jax.tree.map(
-                lambda tar: random.binomial(num_input, p, tar.shape).astype(float),
-                tar_val,
-                is_leaf=u.math.is_quantity
-            )
+            tar_val,
+            is_leaf=u.math.is_quantity
+        )
+        branch2 = jax.tree.map(
+            lambda tar: random.binomial(
+                num_input,
+                p,
+                tar.shape,
+                check_valid=False,
+                dtype=tar.dtype
+            ),
+            tar_val,
+            is_leaf=u.math.is_quantity,
+        )
+
+        inp = jax.tree.map(
+            lambda b1, b2: u.math.where(cond, b1, b2),
+            branch1,
+            branch2,
+            is_leaf=u.math.is_quantity,
         )
 
+        # inp = jax.lax.cond(
+        #     cond,
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.normal(
+        #             a,
+        #             b * p,
+        #             tar.shape,
+        #             key=rand_key,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity
+        #     ),
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.binomial(
+        #             num_input,
+        #             p,
+        #             tar.shape,
+        #             key=rand_key,
+        #             check_valid=False,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity,
+        #     ),
+        #     random.split_key()
+        # )
+
         # update target variable
         target.value = jax.tree.map(
             lambda x: x * weight,
@@ -256,19 +298,62 @@ def poisson_input(
 
     else:
         # generate Poisson input
-        inp = cond(
-            u.math.logical_and(a > 5, b > 5),
-            lambda: jax.tree.map(
-                lambda tar: random.normal(a, b * p, tar[indices].shape),
-                tar_val,
-                is_leaf=u.math.is_quantity
+        branch1 = jax.tree.map(
+            lambda tar: random.normal(
+                a,
+                b * p,
+                tar[indices].shape,
+                dtype=tar.dtype
             ),
-            lambda: jax.tree.map(
-                lambda tar: random.binomial(num_input, p, tar[indices].shape).astype(float),
-                tar_val,
-                is_leaf=u.math.is_quantity
-            )
+            tar_val,
+            is_leaf=u.math.is_quantity
         )
+        branch2 = jax.tree.map(
+            lambda tar: random.binomial(
+                num_input,
+                p,
+                tar[indices].shape,
+                # check_valid=False,
+                dtype=tar.dtype
+            ),
+            tar_val,
+            is_leaf=u.math.is_quantity
+        )
+
+        inp = jax.tree.map(
+            lambda b1, b2: u.math.where(cond, b1, b2),
+            branch1,
+            branch2,
+            is_leaf=u.math.is_quantity,
+        )
+
+        # inp = jax.lax.cond(
+        #     cond,
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.normal(
+        #             a,
+        #             b * p,
+        #             tar[indices].shape,
+        #             key=rand_key,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity
+        #     ),
+        #     lambda rand_key: jax.tree.map(
+        #         lambda tar: random.binomial(
+        #             num_input,
+        #             p,
+        #             tar[indices].shape,
+        #             key=rand_key,
+        #             check_valid=False,
+        #             dtype=tar.dtype
+        #         ),
+        #         tar_val,
+        #         is_leaf=u.math.is_quantity
+        #     ),
+        #     random.split_key()
+        # )
 
         # update target variable
         target.value = jax.tree.map(

brainstate/nn/_interaction/_conv.py

@@ -191,6 +191,7 @@ class _Conv(_BaseConv):
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         w_mask: Optional[Union[ArrayLike, Callable]] = None,
         name: str = None,
+        param_type: type = ParamState,
     ):
         super().__init__(in_size=in_size,
                          out_channels=out_channels,
@@ -215,7 +216,7 @@ class _Conv(_BaseConv):
             params['bias'] = bias
 
         # The weight operation
-        self.weight = ParamState(params)
+        self.weight = param_type(params)
 
         # Evaluate the output shape
         abstract_y = jax.eval_shape(
@@ -346,6 +347,7 @@ class _ScaledWSConv(_BaseConv):
         b_init: Optional[Union[Callable, ArrayLike]] = None,
         w_mask: Optional[Union[ArrayLike, Callable]] = None,
         name: str = None,
+        param_type: type = ParamState,
     ):
         super().__init__(in_size=in_size,
                          out_channels=out_channels,
@@ -379,7 +381,7 @@ class _ScaledWSConv(_BaseConv):
         self.eps = eps
 
         # The weight operation
-        self.weight = ParamState(params)
+        self.weight = param_type(params)
 
         # Evaluate the output shape
         abstract_y = jax.eval_shape(