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

brainstate/event/_csr_test.py

@@ -18,6 +18,9 @@
 import unittest
 
 import brainunit as u
+import jax
+import jax.numpy as jnp
+import numpy as np
 
 import brainstate as bst
 
@@ -88,3 +91,187 @@ class TestCSR(unittest.TestCase):
                 v @ csr
             )
         )
+
+
+def _get_csr(n_pre, n_post, prob):
+    n_conn = int(n_post * prob)
+    indptr = np.arange(n_pre + 1) * n_conn
+    indices = np.random.randint(0, n_post, (n_pre * n_conn,))
+    return indptr, indices
+
+
+def vector_csr(x, w, indices, indptr, shape):
+    homo_w = jnp.size(w) == 1
+    post = jnp.zeros((shape[1],))
+    for i_pre in range(x.shape[0]):
+        ids = indices[indptr[i_pre]: indptr[i_pre + 1]]
+        post = post.at[ids].add(w * x[i_pre] if homo_w else w[indptr[i_pre]: indptr[i_pre + 1]] * x[i_pre])
+    return post
+
+
+def matrix_csr(xs, w, indices, indptr, shape):
+    homo_w = jnp.size(w) == 1
+    post = jnp.zeros((xs.shape[0], shape[1]))
+    for i_pre in range(xs.shape[1]):
+        ids = indices[indptr[i_pre]: indptr[i_pre + 1]]
+        post = post.at[:, ids].add(
+            w * xs[:, i_pre: i_pre + 1]
+            if homo_w else
+            (w[indptr[i_pre]: indptr[i_pre + 1]] * xs[:, i_pre: i_pre + 1])
+        )
+    return post
+
+
+def csr_vector(x, w, indices, indptr, shape):
+    homo_w = jnp.size(w) == 1
+    out = jnp.zeros([shape[0]])
+    for i in range(shape[0]):
+        ids = indices[indptr[i]: indptr[i + 1]]
+        ws = w if homo_w else w[indptr[i]: indptr[i + 1]]
+        out = out.at[i].set(jnp.sum(x[ids] * ws))
+    return out
+
+
+def csr_matrix(xs, w, indices, indptr, shape):
+    # CSR @ matrix
+    homo_w = jnp.size(w) == 1
+    out = jnp.zeros([shape[0], xs.shape[1]])
+    for i in range(shape[0]):
+        ids = indices[indptr[i]: indptr[i + 1]]
+        ws = w if homo_w else jnp.expand_dims(w[indptr[i]: indptr[i + 1]], axis=1)
+        out = out.at[i].set(jnp.sum(xs[ids] * ws, axis=0))
+    return out
+
+
+class TestVectorCSR(unittest.TestCase):
+    def test_vector_csr(self, ):
+        m, n = 20, 40
+        x = bst.random.rand(m) < 0.1
+        indptr, indices = _get_csr(m, n, 0.1)
+
+        for homo_w in [True, False]:
+            print(f'homo_w = {homo_w}')
+            data = 1.5 if homo_w else bst.init.Normal()(indices.shape)
+            csr = bst.event.CSR([data, indices, indptr], shape=(m, n))
+            y = x @ csr
+            y2 = vector_csr(x, csr.data, indices, indptr, [m, n])
+            self.assertTrue(jnp.allclose(y, y2))
+
+    def test_vector_csr_vmap_vector(self):
+        n_batch, m, n = 10, 20, 40
+        xs = bst.random.rand(n_batch, m) < 0.1
+        indptr, indices = _get_csr(m, n, 0.1)
+
+        for homo_w in [True, False]:
+            data = 1.5 if homo_w else bst.init.Normal()(indices.shape)
+            csr = bst.event.CSR([data, indices, indptr], shape=(m, n))
+            y = jax.vmap(lambda x: x @ csr)(xs)
+            y2 = jax.vmap(lambda x: vector_csr(x, csr.data, indices, indptr, [m, n]))(xs)
+            self.assertTrue(jnp.allclose(y, y2))
+
+
+class TestMatrixCSR(unittest.TestCase):
+    def test_matrix_csr(self):
+        k, m, n = 10, 20, 40
+        x = bst.random.rand(k, m) < 0.1
+        indptr, indices = _get_csr(m, n, 0.1)
+
+        for homo_w in [True, False]:
+            data = 1.5 if homo_w else bst.init.Normal()(indices.shape)
+            csr = bst.event.CSR([data, indices, indptr], shape=(m, n))
+            y = x @ csr
+            y2 = matrix_csr(x, csr.data, indices, indptr, [m, n])
+            self.assertTrue(jnp.allclose(y, y2))
+
+
+class TestCSRVector(unittest.TestCase):
+    def test_csr_vector(self):
+        m, n = 20, 40
+        v = bst.random.rand(n) < 0.1
+        indptr, indices = _get_csr(m, n, 0.1)
+
+        for homo_w in [True, False]:
+            data = 1.5 if homo_w else bst.init.Normal()(indices.shape)
+            csr = bst.event.CSR([data, indices, indptr], shape=(m, n))
+            y = csr @ v
+            y2 = csr_vector(v, csr.data, indices, indptr, [m, n])
+            self.assertTrue(jnp.allclose(y, y2))
+
+
+class TestCSRMatrix(unittest.TestCase):
+    def test_csr_matrix(self):
+        m, n, k = 20, 40, 10
+        matrix = bst.random.rand(n, k) < 0.1
+        indptr, indices = _get_csr(m, n, 0.1)
+
+        for homo_w in [True, False]:
+            data = 1.5 if homo_w else bst.init.Normal()(indices.shape)
+            csr = bst.event.CSR([data, indices, indptr], shape=(m, n))
+            y = csr @ matrix
+            y2 = csr_matrix(matrix, csr.data, indices, indptr, [m, n])
+            self.assertTrue(jnp.allclose(y, y2))
+
+    # @parameterized.product(
+    #     bool_x=[True, False],
+    #     homo_w=[True, False]
+    # )
+    # def test_vjp(self, bool_x, homo_w):
+    #     n_in = 20
+    #     n_out = 30
+    #     if bool_x:
+    #         x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
+    #     else:
+    #         x = bst.random.rand(n_in)
+    #
+    #     indptr, indices = _get_csr(n_in, n_out, 0.1)
+    #     fn = bst.event.CSRLinear(n_in, n_out, indptr, indices, 1.5 if homo_w else bst.init.Normal())
+    #     w = fn.weight.value
+    #
+    #     def f(x, w):
+    #         fn.weight.value = w
+    #         return fn(x).sum()
+    #
+    #     r = jax.grad(f, argnums=(0, 1))(x, w)
+    #
+    #     # -------------------
+    #     # TRUE gradients
+    #
+    #     def f2(x, w):
+    #         return true_fn(x, w, indices, indptr, n_out).sum()
+    #
+    #     r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+    #     self.assertTrue(jnp.allclose(r[0], r2[0]))
+    #     self.assertTrue(jnp.allclose(r[1], r2[1]))
+    #
+    # @parameterized.product(
+    #     bool_x=[True, False],
+    #     homo_w=[True, False]
+    # )
+    # def test_jvp(self, bool_x, homo_w):
+    #     n_in = 20
+    #     n_out = 30
+    #     if bool_x:
+    #         x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
+    #     else:
+    #         x = bst.random.rand(n_in)
+    #
+    #     indptr, indices = _get_csr(n_in, n_out, 0.1)
+    #     fn = bst.event.CSRLinear(n_in, n_out, indptr, indices,
+    #                              1.5 if homo_w else bst.init.Normal(), grad_mode='jvp')
+    #     w = fn.weight.value
+    #
+    #     def f(x, w):
+    #         fn.weight.value = w
+    #         return fn(x)
+    #
+    #     o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+    #
+    #     # -------------------
+    #     # TRUE gradients
+    #
+    #     def f2(x, w):
+    #         return true_fn(x, w, indices, indptr, n_out)
+    #
+    #     o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+    #     self.assertTrue(jnp.allclose(r1, r2))
+    #     self.assertTrue(jnp.allclose(o1, o2))

brainstate/event/_xla_custom_op.py

@@ -9,7 +9,6 @@ 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 jaxlib.hlo_helpers import custom_call
@@ -19,6 +18,11 @@ if jax.__version_info__ < (0, 4, 35):
 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__ = [
@@ -164,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
 ):
@@ -205,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