brainstate 0.1.0.post20241209__py2.py3-none-any.whl → 0.1.0.post20241219__py2.py3-none-any.whl

brainstate/event/_csr_mv.py

@@ -58,7 +58,6 @@ class CSRLinear(Module):
         indices: ArrayLike,
         weight: Union[Callable, ArrayLike],
         name: Optional[str] = None,
-        grad_mode: str = 'vjp'
     ):
         super().__init__(name=name)
 
@@ -68,17 +67,13 @@ class CSRLinear(Module):
         self.n_pre = self.in_size[-1]
         self.n_post = self.out_size[-1]
 
-        # gradient mode
-        assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
-        self.grad_mode = grad_mode
-
         # CSR data structure
-        indptr = jnp.asarray(indptr)
-        indices = jnp.asarray(indices)
-        assert indptr.ndim == 1, f"indptr must be 1D. Got: {indptr.ndim}"
-        assert indices.ndim == 1, f"indices must be 1D. Got: {indices.ndim}"
-        assert indptr.size == self.n_pre + 1, f"indptr must have size {self.n_pre + 1}. Got: {indptr.size}"
         with jax.ensure_compile_time_eval():
+            indptr = jnp.asarray(indptr)
+            indices = jnp.asarray(indices)
+            assert indptr.ndim == 1, f"indptr must be 1D. Got: {indptr.ndim}"
+            assert indices.ndim == 1, f"indices must be 1D. Got: {indices.ndim}"
+            assert indptr.size == self.n_pre + 1, f"indptr must have size {self.n_pre + 1}. Got: {indptr.size}"
             self.indptr = u.math.asarray(indptr)
             self.indices = u.math.asarray(indices)
 
@@ -101,21 +96,13 @@ class CSRLinear(Module):
         device_kind = jax.devices()[0].platform  # spk.device.device_kind
 
         # CPU implementation
-        if device_kind == 'cpu':
-            return cpu_event_csr(
-                u.math.asarray(spk),
-                self.indptr,
-                self.indices,
-                u.math.asarray(weight),
-                n_post=self.n_post, grad_mode=self.grad_mode
-            )
-
-        # GPU/TPU implementation
-        elif device_kind in ['gpu', 'tpu']:
-            raise NotImplementedError()
-
-        else:
-            raise ValueError(f"Unsupported device: {device_kind}")
+        return cpu_event_csr(
+            u.math.asarray(spk),
+            self.indptr,
+            self.indices,
+            u.math.asarray(weight),
+            n_post=self.n_post,
+        )
 
 
 @set_module_as('brainstate.event')

brainstate/event/_csr_mv_test.py

@@ -40,79 +40,79 @@ def true_fn(x, w, indices, indptr, n_out):
     return post
 
 
-class TestFixedProbCSR(parameterized.TestCase):
-    @parameterized.product(
-        homo_w=[True, False],
-    )
-    def test1(self, homo_w):
-        x = bst.random.rand(20) < 0.1
-        indptr, indices = _get_csr(20, 40, 0.1)
-        m = bst.event.CSRLinear(20, 40, indptr, indices, 1.5 if homo_w else bst.init.Normal())
-        y = m(x)
-        y2 = true_fn(x, m.weight.value, indices, indptr, 40)
-        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))
+# class TestFixedProbCSR(parameterized.TestCase):
+#     @parameterized.product(
+#         homo_w=[True, False],
+#     )
+#     def test1(self, homo_w):
+#         x = bst.random.rand(20) < 0.1
+#         indptr, indices = _get_csr(20, 40, 0.1)
+#         m = bst.event.CSRLinear(20, 40, indptr, indices, 1.5 if homo_w else bst.init.Normal())
+#         y = m(x)
+#         y2 = true_fn(x, m.weight.value, indices, indptr, 40)
+#         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/_csr_test.py

@@ -0,0 +1,90 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# -*- coding: utf-8 -*-
+
+
+import unittest
+
+import brainunit as u
+
+import brainstate as bst
+
+
+class TestCSR(unittest.TestCase):
+    def test_event_homo_bool(self):
+        for dat in [1., 2., 3.]:
+            mask = (bst.random.rand(10, 20) < 0.1).astype(float) * dat
+            csr = u.sparse.CSR.fromdense(mask)
+            csr = bst.event.CSR((dat, csr.indices, csr.indptr), shape=mask.shape)
+
+            v = bst.random.rand(20) < 0.5
+            self.assertTrue(
+                u.math.allclose(
+                    mask.astype(float) @ v.astype(float),
+                    csr @ v
+                )
+            )
+
+            v = bst.random.rand(10) < 0.5
+            self.assertTrue(
+                u.math.allclose(
+                    v.astype(float) @ mask.astype(float),
+                    v @ csr
+                )
+            )
+
+    def test_event_homo_heter(self):
+        mat = bst.random.rand(10, 20)
+        mask = (bst.random.rand(10, 20) < 0.1) * mat
+        csr = u.sparse.CSR.fromdense(mask)
+        csr = bst.event.CSR((csr.data, csr.indices, csr.indptr), shape=mask.shape)
+
+        v = bst.random.rand(20) < 0.5
+        self.assertTrue(
+            u.math.allclose(
+                mask.astype(float) @ v.astype(float),
+                csr @ v
+            )
+        )
+
+        v = bst.random.rand(10) < 0.5
+        self.assertTrue(
+            u.math.allclose(
+                v.astype(float) @ mask.astype(float),
+                v @ csr
+            )
+        )
+
+    def test_event_heter_float_as_bool(self):
+        mat = bst.random.rand(10, 20)
+        mask = (mat < 0.1).astype(float) * mat
+        csr = u.sparse.CSR.fromdense(mask)
+        csr = bst.event.CSR((csr.data, csr.indices, csr.indptr), shape=mask.shape)
+
+        v = (bst.random.rand(20) < 0.5).astype(float)
+        self.assertTrue(
+            u.math.allclose(
+                mask.astype(float) @ v.astype(float),
+                csr @ v
+            )
+        )
+
+        v = (bst.random.rand(10) < 0.5).astype(float)
+        self.assertTrue(
+            u.math.allclose(
+                v.astype(float) @ mask.astype(float),
+                v @ csr
+            )
+        )

brainstate/event/_fixedprob_mv.py

@@ -85,44 +85,52 @@ class FixedProb(Module):
         self.in_size = in_size
         self.out_size = out_size
         self.n_conn = int(self.out_size[-1] * prob)
-        if self.n_conn < 1:
-            raise ValueError(f"The number of connections must be at least 1. "
-                             f"Got: int({self.out_size[-1]} * {prob}) = {self.n_conn}")
         self.float_as_event = float_as_event
         self.block_size = block_size
 
-        # indices of post connected neurons
-        with jax.ensure_compile_time_eval():
-            if allow_multi_conn:
-                rng = np.random.RandomState(seed)
-                self.indices = rng.randint(0, self.out_size[-1], size=(self.in_size[-1], self.n_conn))
-            else:
-                rng = RandomState(seed)
+        if self.n_conn > 1:
+            # indices of post connected neurons
+            with jax.ensure_compile_time_eval():
+                if allow_multi_conn:
+                    rng = np.random.RandomState(seed)
+                    self.indices = rng.randint(0, self.out_size[-1], size=(self.in_size[-1], self.n_conn))
+                else:
+                    rng = RandomState(seed)
 
-                @vmap(rngs=rng)
-                def rand_indices(key):
-                    rng.set_key(key)
-                    return rng.choice(self.out_size[-1], size=(self.n_conn,), replace=False)
+                    @vmap(rngs=rng)
+                    def rand_indices(key):
+                        rng.set_key(key)
+                        return rng.choice(self.out_size[-1], size=(self.n_conn,), replace=False)
 
-                self.indices = rand_indices(rng.split_key(self.in_size[-1]))
-            self.indices = u.math.asarray(self.indices)
+                    self.indices = rand_indices(rng.split_key(self.in_size[-1]))
+                self.indices = u.math.asarray(self.indices)
 
         # maximum synaptic conductance
         weight = param(weight, (self.in_size[-1], self.n_conn), allow_none=False)
         self.weight = ParamState(weight)
 
     def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-        return event_fixed_prob(
-            spk,
-            self.weight.value,
-            self.indices,
-            n_post=self.out_size[-1],
-            block_size=self.block_size,
-            float_as_event=self.float_as_event
-        )
+        if self.n_conn > 1:
+            return event_fixed_prob(
+                spk,
+                self.weight.value,
+                self.indices,
+                n_post=self.out_size[-1],
+                block_size=self.block_size,
+                float_as_event=self.float_as_event
+            )
+        else:
+            weight = self.weight.value
+            unit = u.get_unit(weight)
+            r = jnp.zeros(spk.shape[:-1] + (self.out_size[-1],), dtype=weight.dtype)
+            return u.maybe_decimal(u.Quantity(r, unit=unit))
 
 
-def event_fixed_prob(spk, weight, indices, *, n_post, block_size, float_as_event):
+def event_fixed_prob(
+    spk, weight, indices,
+    *,
+    n_post, block_size, float_as_event
+):
     """
     The FixedProb module implements a fixed probability connection with CSR sparse data structure.
 
@@ -374,7 +382,11 @@ def gpu_kernel_generator(
                 kernel(spikes, indices, weight, jnp.zeros(n_post, dtype=weight_info.dtype)))
 
 
-def jvp_spikes(spk_dot, spikes, weights, indices, *, n_post, block_size, **kwargs):
+def jvp_spikes(
+    spk_dot, spikes, weights, indices,
+    *,
+    n_post, block_size, **kwargs
+):
     return ellmv_p_call(
         spk_dot,
         weights,
@@ -384,7 +396,11 @@ def jvp_spikes(spk_dot, spikes, weights, indices, *, n_post, block_size, **kwarg
     )
 
 
-def jvp_weights(w_dot, spikes, weights, indices, *, float_as_event, block_size, n_post, **kwargs):
+def jvp_weights(
+    w_dot, spikes, weights, indices,
+    *,
+    float_as_event, block_size, n_post, **kwargs
+):
     return event_ellmv_p_call(
         spikes,
         w_dot,
@@ -457,14 +473,18 @@ def transpose_rule(
 
 event_ellmv_p = XLACustomOp(
     'event_ell_mv',
-    cpu_kernel_generator=cpu_kernel_generator,
-    gpu_kernel_generator=gpu_kernel_generator,
+    cpu_kernel_or_generator=cpu_kernel_generator,
+    gpu_kernel_or_generator=gpu_kernel_generator,
 )
 event_ellmv_p.defjvp(jvp_spikes, jvp_weights, None)
 event_ellmv_p.def_transpose_rule(transpose_rule)
 
 
-def event_ellmv_p_call(spikes, weights, indices, *, n_post, block_size, float_as_event):
+def event_ellmv_p_call(
+    spikes, weights, indices,
+    *,
+    n_post, block_size, float_as_event
+):
     n_conn = indices.shape[1]
     if block_size is None:
         if n_conn <= 16:
@@ -673,8 +693,8 @@ def transpose_rule_no_spk(
 
 ellmv_p = XLACustomOp(
     'ell_mv',
-    cpu_kernel_generator=ell_cpu_kernel_generator,
-    gpu_kernel_generator=ell_gpu_kernel_generator,
+    cpu_kernel_or_generator=ell_cpu_kernel_generator,
+    gpu_kernel_or_generator=ell_gpu_kernel_generator,
 )
 ellmv_p.defjvp(jvp_spikes, jvp_weights_no_spk, None)
 ellmv_p.def_transpose_rule(transpose_rule_no_spk)

brainstate/event/_linear_mv.py

@@ -334,8 +334,8 @@ def transpose_rule(ct, spikes, weights, *, float_as_event, **kwargs):
 
 event_linear_p = XLACustomOp(
     'event_linear',
-    cpu_kernel_generator=cpu_kernel_generator,
-    gpu_kernel_generator=gpu_kernel_generator,
+    cpu_kernel_or_generator=cpu_kernel_generator,
+    gpu_kernel_or_generator=gpu_kernel_generator,
 )
 event_linear_p.defjvp(jvp_spikes, jvp_weights)
 event_linear_p.def_transpose_rule(transpose_rule)

brainstate/event/_xla_custom_op.py

@@ -180,8 +180,8 @@ class XLACustomOp:
     """Creating a XLA custom call operator.
 
     Args:
-      cpu_kernel_generator: Callable. The function defines the computation on CPU backend.
-      gpu_kernel_generator: Callable. The function defines the computation on GPU backend.
+      cpu_kernel_or_generator: Callable. The function defines the computation on CPU backend.
+      gpu_kernel_or_generator: Callable. The function defines the computation on GPU backend.
       batching_translation: Callable. The batching translation rule of JAX.
       jvp_translation: Callable. The JVP translation rule of JAX.
       transpose_translation: Callable. The transpose translation rule of JAX.
@@ -191,15 +191,12 @@ class XLACustomOp:
     def __init__(
         self,
         name: str,
-        cpu_kernel_generator: Callable,
-        gpu_kernel_generator: Callable = None,
+        cpu_kernel_or_generator: Callable,
+        gpu_kernel_or_generator: Callable = None,
         batching_translation: Callable = None,
         jvp_translation: Callable = None,
         transpose_translation: Callable = None,
     ):
-        # set cpu_kernel and gpu_kernel
-        self.cpu_kernel = cpu_kernel_generator
-
         # primitive
         self.primitive = jax.core.Primitive(name)
         self.primitive.multiple_results = True
@@ -209,10 +206,10 @@ class XLACustomOp:
         self.primitive.def_abstract_eval(self._abstract_eval)
 
         # cpu kernel
-        if cpu_kernel_generator is not None:
-            self.def_cpu_kernel(cpu_kernel_generator)
-        if gpu_kernel_generator is not None:
-            self.def_gpu_kernel(gpu_kernel_generator)
+        if cpu_kernel_or_generator is not None:
+            self.def_cpu_kernel(cpu_kernel_or_generator)
+        if gpu_kernel_or_generator is not None:
+            self.def_gpu_kernel(gpu_kernel_or_generator)
 
         # batching rule
         if batching_translation is not None:

brainstate/graph/_graph_node.py

@@ -61,6 +61,9 @@ class Node(PrettyRepr, metaclass=GraphNodeMeta):
     - Deepcopy the node.
 
     """
+
+    graph_invisible_attrs = ()
+
     if TYPE_CHECKING:
         _trace_state: StateJaxTracer
 
@@ -170,7 +173,13 @@ def _to_shape_dtype(value):
 def _node_flatten(
     node: Node
 ) -> Tuple[Tuple[Tuple[str, Any], ...], Tuple[Type]]:
-    nodes = sorted((key, value) for key, value in vars(node).items() if key != '_trace_state')
+    # graph_invisible_attrs = getattr(node, 'graph_invisible_attrs', ())
+    # graph_invisible_attrs = tuple(graph_invisible_attrs) + ('_trace_state',)
+    graph_invisible_attrs = ('_trace_state',)
+    nodes = sorted(
+        (key, value) for key, value in vars(node).items()
+        if (key not in graph_invisible_attrs)
+    )
     return nodes, (type(node),)
 
 

brainstate/graph/_graph_operation.py

@@ -608,9 +608,9 @@ def _get_children(graph_def, state_mapping, index_ref, index_ref_cache):
                         if isinstance(value, TreefyState):
                             variable.update_from_ref(value)
                         elif isinstance(value, State):
-                             if value._been_writen:
+                            if value._been_writen:
                                 variable.write_value(value.value)
-                             else:
+                            else:
                                 variable.restore_value(value.value)
                         else:
                             raise ValueError(f'Expected a State type for {key!r}, but got {type(value)}.')
@@ -1600,10 +1600,12 @@ def iter_leaf(
             visited_.add(id(node_))
             node_dict = _get_node_impl(node_).node_dict(node_)
             for key, value in node_dict.items():
-                yield from _iter_graph_leaf(value,
-                                            visited_,
-                                            (*path_parts_, key),
-                                            level_ + 1 if _is_graph_node(value) else level_)
+                yield from _iter_graph_leaf(
+                    value,
+                    visited_,
+                    (*path_parts_, key),
+                    level_ + 1 if _is_graph_node(value) else level_
+                )
         else:
             if level_ >= allowed_hierarchy[0]:
                 yield path_parts_, node_