brainstate 0.1.10__py2.py3-none-any.whl → 0.2.1__py2.py3-none-any.whl

brainstate/nn/{_fixedprob_test.py → _event_fixedprob_test.py}

@@ -1,114 +1,115 @@
-# 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.
-# ==============================================================================
-
-
-import jax.numpy
-import jax.numpy as jnp
-import pytest
-
-import brainstate
-
-
-class TestFixedProbCSR:
-    @pytest.mark.parametrize('allow_multi_conn', [True, False, ])
-    def test1(self, allow_multi_conn):
-        x = brainstate.random.rand(20) < 0.1
-        # x = brainstate.random.rand(20)
-        m = brainstate.nn.EventFixedProb(20, 40, 0.1, 1.0, seed=123, allow_multi_conn=allow_multi_conn)
-        y = m(x)
-        print(y)
-
-        m2 = brainstate.nn.EventFixedProb(20, 40, 0.1, brainstate.init.KaimingUniform(), seed=123)
-        print(m2(x))
-
-    def test_grad_bool(self):
-        n_in = 20
-        n_out = 30
-        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
-        fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, brainstate.init.KaimingUniform(), seed=123)
-
-        def f(x):
-            return fn(x).sum()
-
-        print(jax.grad(f)(x))
-
-    @pytest.mark.parametrize('homo_w', [True, False])
-    def test_vjp(self, homo_w):
-        n_in = 20
-        n_out = 30
-        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
-
-        if homo_w:
-            fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, 1.5, seed=123)
-        else:
-            fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, brainstate.init.KaimingUniform(), seed=123)
-        w = fn.weight.value
-
-        def f(x, w):
-            fn.weight.value = w
-            return fn(x).sum()
-
-        r = brainstate.augment.grad(f, argnums=(0, 1))(x, w)
-
-        # -------------------
-        # TRUE gradients
-
-        def true_fn(x, w, indices, n_post):
-            post = jnp.zeros((n_post,))
-            for i in range(n_in):
-                post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
-            return post
-
-        def f2(x, w):
-            return true_fn(x, w, fn.conn.indices, n_out).sum()
-
-        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
-        assert (jnp.allclose(r[0], r2[0]))
-        assert (jnp.allclose(r[1], r2[1]))
-
-    @pytest.mark.parametrize('homo_w', [True, False])
-    def test_jvp(self, homo_w):
-        n_in = 20
-        n_out = 30
-        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
-
-        fn = brainstate.nn.EventFixedProb(
-            n_in, n_out, 0.1, 1.5 if homo_w else brainstate.init.KaimingUniform(),
-            seed=123,
-        )
-        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 true_fn(x, w, indices, n_post):
-            post = jnp.zeros((n_post,))
-            for i in range(n_in):
-                post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
-            return post
-
-        def f2(x, w):
-            return true_fn(x, w, fn.conn.indices, n_out)
-
-        o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-        assert (jnp.allclose(o1, o2))
-        # assert jnp.allclose(r1, r2), f'r1={r1}, r2={r2}'
-        assert (jnp.allclose(r1, r2, rtol=1e-4, atol=1e-4))
+# Copyright 2024 BrainX 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.
+# ==============================================================================
+
+
+import jax.numpy
+import jax.numpy as jnp
+import pytest
+
+import brainstate
+import braintools
+
+
+class TestFixedProbCSR:
+    @pytest.mark.parametrize('allow_multi_conn', [True, False, ])
+    def test1(self, allow_multi_conn):
+        x = brainstate.random.rand(20) < 0.1
+        # x = brainstate.random.rand(20)
+        m = brainstate.nn.EventFixedProb(20, 40, 0.1, 1.0, seed=123, allow_multi_conn=allow_multi_conn)
+        y = m(x)
+        print(y)
+
+        m2 = brainstate.nn.EventFixedProb(20, 40, 0.1, braintools.init.KaimingUniform(), seed=123)
+        print(m2(x))
+
+    def test_grad_bool(self):
+        n_in = 20
+        n_out = 30
+        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+        fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, braintools.init.KaimingUniform(), seed=123)
+
+        def f(x):
+            return fn(x).sum()
+
+        print(jax.grad(f)(x))
+
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test_vjp(self, homo_w):
+        n_in = 20
+        n_out = 30
+        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+
+        if homo_w:
+            fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, 1.5, seed=123)
+        else:
+            fn = brainstate.nn.EventFixedProb(n_in, n_out, 0.1, braintools.init.KaimingUniform(), seed=123)
+        w = fn.weight.value
+
+        def f(x, w):
+            fn.weight.value = w
+            return fn(x).sum()
+
+        r = brainstate.augment.grad(f, argnums=(0, 1))(x, w)
+
+        # -------------------
+        # TRUE gradients
+
+        def true_fn(x, w, indices, n_post):
+            post = jnp.zeros((n_post,))
+            for i in range(n_in):
+                post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
+            return post
+
+        def f2(x, w):
+            return true_fn(x, w, fn.conn.indices, n_out).sum()
+
+        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+        assert (jnp.allclose(r[0], r2[0]))
+        assert (jnp.allclose(r[1], r2[1]))
+
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test_jvp(self, homo_w):
+        n_in = 20
+        n_out = 30
+        x = jax.numpy.asarray(brainstate.random.rand(n_in) < 0.3, dtype=float)
+
+        fn = brainstate.nn.EventFixedProb(
+            n_in, n_out, 0.1, 1.5 if homo_w else braintools.init.KaimingUniform(),
+            seed=123,
+        )
+        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 true_fn(x, w, indices, n_post):
+            post = jnp.zeros((n_post,))
+            for i in range(n_in):
+                post = post.at[indices[i]].add(w * x[i] if homo_w else w[i] * x[i])
+            return post
+
+        def f2(x, w):
+            return true_fn(x, w, fn.conn.indices, n_out)
+
+        o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+        assert (jnp.allclose(o1, o2))
+        # assert jnp.allclose(r1, r2), f'r1={r1}, r2={r2}'
+        assert (jnp.allclose(r1, r2, rtol=1e-4, atol=1e-4))

brainstate/nn/{_linear_mv.py → _event_linear.py}

@@ -1,83 +1,83 @@
-# 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.
-# ==============================================================================
-
-from typing import Union, Callable, Optional
-
-import brainevent
-import brainunit as u
-import jax
-
-from brainstate import init
-from brainstate._state import ParamState
-from brainstate.typing import Size, ArrayLike
-from ._module import Module
-
-__all__ = [
-    'EventLinear',
-]
-
-
-class EventLinear(Module):
-    """
-
-    Parameters
-    ----------
-    in_size : Size
-        Number of pre-synaptic neurons, i.e., input size.
-    out_size : Size
-        Number of post-synaptic neurons, i.e., output size.
-    weight : float or callable or jax.Array or brainunit.Quantity
-        Maximum synaptic conductance.
-    block_size : int, optional
-        Block size for parallel computation.
-    float_as_event : bool, optional
-        Whether to treat float as event.
-    name : str, optional
-        Name of the module.
-    """
-
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        out_size: Size,
-        weight: Union[Callable, ArrayLike],
-        float_as_event: bool = True,
-        block_size: int = 64,
-        name: Optional[str] = None,
-        param_type: type = ParamState,
-    ):
-        super().__init__(name=name)
-
-        # network parameters
-        self.in_size = in_size
-        self.out_size = out_size
-        self.float_as_event = float_as_event
-        self.block_size = block_size
-
-        # maximum synaptic conductance
-        weight = init.param(weight, (self.in_size[-1], self.out_size[-1]), allow_none=False)
-        self.weight = param_type(weight)
-
-    def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-        weight = self.weight.value
-        if u.math.size(weight) == 1:
-            return u.math.ones(self.out_size) * (u.math.sum(spk) * weight)
-
-        if self.float_as_event:
-            return brainevent.EventArray(spk) @ weight
-        else:
-            return spk @ weight
+# Copyright 2024 BrainX 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.
+# ==============================================================================
+
+from typing import Union, Callable, Optional
+
+import brainevent
+import brainunit as u
+import jax
+
+from brainstate._state import ParamState
+from brainstate.typing import Size, ArrayLike
+from . import init as init
+from ._module import Module
+
+__all__ = [
+    'EventLinear',
+]
+
+
+class EventLinear(Module):
+    """
+
+    Parameters
+    ----------
+    in_size : Size
+        Number of pre-synaptic neurons, i.e., input size.
+    out_size : Size
+        Number of post-synaptic neurons, i.e., output size.
+    weight : float or callable or jax.Array or brainunit.Quantity
+        Maximum synaptic conductance.
+    block_size : int, optional
+        Block size for parallel computation.
+    float_as_event : bool, optional
+        Whether to treat float as event.
+    name : str, optional
+        Name of the module.
+    """
+
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        out_size: Size,
+        weight: Union[Callable, ArrayLike],
+        float_as_event: bool = True,
+        block_size: int = 64,
+        name: Optional[str] = None,
+        param_type: type = ParamState,
+    ):
+        super().__init__(name=name)
+
+        # network parameters
+        self.in_size = in_size
+        self.out_size = out_size
+        self.float_as_event = float_as_event
+        self.block_size = block_size
+
+        # maximum synaptic conductance
+        weight = init.param(weight, (self.in_size[-1], self.out_size[-1]), allow_none=False)
+        self.weight = param_type(weight)
+
+    def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
+        weight = self.weight.value
+        if u.math.size(weight) == 1:
+            return u.math.ones(self.out_size) * (u.math.sum(spk) * weight)
+
+        if self.float_as_event:
+            return brainevent.EventArray(spk) @ weight
+        else:
+            return spk @ weight

brainstate/nn/{_linear_mv_test.py → _event_linear_test.py}

@@ -1,120 +1,121 @@
-# 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.
-# ==============================================================================
-
-
-import jax
-import jax.numpy as jnp
-import pytest
-
-import brainstate
-
-
-class TestEventLinear:
-    @pytest.mark.parametrize('bool_x', [True, False])
-    @pytest.mark.parametrize('homo_w', [True, False])
-    def test1(self, homo_w, bool_x):
-        x = brainstate.random.rand(20) < 0.1
-        if not bool_x:
-            x = jnp.asarray(x, dtype=float)
-        m = brainstate.nn.EventLinear(
-            20, 40,
-            1.5 if homo_w else brainstate.init.KaimingUniform(),
-            float_as_event=bool_x
-        )
-        y = m(x)
-        print(y)
-
-        assert (jnp.allclose(y, (x.sum() * m.weight.value) if homo_w else (x @ m.weight.value)))
-
-    def test_grad_bool(self):
-        n_in = 20
-        n_out = 30
-        x = brainstate.random.rand(n_in) < 0.3
-        fn = brainstate.nn.EventLinear(n_in, n_out, brainstate.init.KaimingUniform())
-
-        with pytest.raises(TypeError):
-            print(jax.grad(lambda x: fn(x).sum())(x))
-
-    @pytest.mark.parametrize('bool_x', [True, False])
-    @pytest.mark.parametrize('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(brainstate.random.rand(n_in) < 0.3, dtype=float)
-        else:
-            x = brainstate.random.rand(n_in)
-
-        fn = brainstate.nn.EventLinear(
-            n_in,
-            n_out,
-            1.5 if homo_w else brainstate.init.KaimingUniform(),
-            float_as_event=bool_x
-        )
-        w = fn.weight.value
-
-        def f(x, w):
-            fn.weight.value = w
-            return fn(x).sum()
-
-        r1 = jax.grad(f, argnums=(0, 1))(x, w)
-
-        # -------------------
-        # TRUE gradients
-
-        def f2(x, w):
-            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
-            return y.sum()
-
-        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
-        assert (jnp.allclose(r1[0], r2[0]))
-
-        if not jnp.allclose(r1[1], r2[1]):
-            print(r1[1] - r2[1])
-
-        assert (jnp.allclose(r1[1], r2[1]))
-
-    @pytest.mark.parametrize('bool_x', [True, False])
-    @pytest.mark.parametrize('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(brainstate.random.rand(n_in) < 0.3, dtype=float)
-        else:
-            x = brainstate.random.rand(n_in)
-
-        fn = brainstate.nn.EventLinear(
-            n_in, n_out, 1.5 if homo_w else brainstate.init.KaimingUniform(),
-            float_as_event=bool_x
-        )
-        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):
-            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
-            return y
-
-        o2, r2 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-        assert (jnp.allclose(o1, o2))
-        assert (jnp.allclose(r1, r2))
+# Copyright 2024 BrainX 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.
+# ==============================================================================
+
+
+import jax
+import jax.numpy as jnp
+import pytest
+
+import braintools
+import brainstate
+
+
+class TestEventLinear:
+    @pytest.mark.parametrize('bool_x', [True, False])
+    @pytest.mark.parametrize('homo_w', [True, False])
+    def test1(self, homo_w, bool_x):
+        x = brainstate.random.rand(20) < 0.1
+        if not bool_x:
+            x = jnp.asarray(x, dtype=float)
+        m = brainstate.nn.EventLinear(
+            20, 40,
+            1.5 if homo_w else braintools.init.KaimingUniform(),
+            float_as_event=bool_x
+        )
+        y = m(x)
+        print(y)
+
+        assert (jnp.allclose(y, (x.sum() * m.weight.value) if homo_w else (x @ m.weight.value)))
+
+    def test_grad_bool(self):
+        n_in = 20
+        n_out = 30
+        x = brainstate.random.rand(n_in) < 0.3
+        fn = brainstate.nn.EventLinear(n_in, n_out, braintools.init.KaimingUniform())
+
+        with pytest.raises(TypeError):
+            print(jax.grad(lambda x: fn(x).sum())(x))
+
+    @pytest.mark.parametrize('bool_x', [True, False])
+    @pytest.mark.parametrize('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(brainstate.random.rand(n_in) < 0.3, dtype=float)
+        else:
+            x = brainstate.random.rand(n_in)
+
+        fn = brainstate.nn.EventLinear(
+            n_in,
+            n_out,
+            1.5 if homo_w else braintools.init.KaimingUniform(),
+            float_as_event=bool_x
+        )
+        w = fn.weight.value
+
+        def f(x, w):
+            fn.weight.value = w
+            return fn(x).sum()
+
+        r1 = jax.grad(f, argnums=(0, 1))(x, w)
+
+        # -------------------
+        # TRUE gradients
+
+        def f2(x, w):
+            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
+            return y.sum()
+
+        r2 = jax.grad(f2, argnums=(0, 1))(x, w)
+        assert (jnp.allclose(r1[0], r2[0]))
+
+        if not jnp.allclose(r1[1], r2[1]):
+            print(r1[1] - r2[1])
+
+        assert (jnp.allclose(r1[1], r2[1]))
+
+    @pytest.mark.parametrize('bool_x', [True, False])
+    @pytest.mark.parametrize('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(brainstate.random.rand(n_in) < 0.3, dtype=float)
+        else:
+            x = brainstate.random.rand(n_in)
+
+        fn = brainstate.nn.EventLinear(
+            n_in, n_out, 1.5 if homo_w else braintools.init.KaimingUniform(),
+            float_as_event=bool_x
+        )
+        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):
+            y = (x @ (jnp.ones([n_in, n_out]) * w)) if homo_w else (x @ w)
+            return y
+
+        o2, r2 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
+        assert (jnp.allclose(o1, o2))
+        assert (jnp.allclose(r1, r2))