brainstate 0.1.7__py2.py3-none-any.whl → 0.1.9__py2.py3-none-any.whl

brainstate/nn/_linear_mv.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 brainunit as u
-import jax
-
-from brainstate import init
-import brainevent
-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 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

brainstate/nn/_linear_mv_test.py

@@ -1,120 +1,120 @@
-# 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 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))

brainstate/nn/_linear_test.py

@@ -1,107 +1,107 @@
-# 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 unittest
-
-import brainunit as u
-from absl.testing import parameterized
-
-import brainstate
-
-
-class TestDense(parameterized.TestCase):
-    @parameterized.product(
-        size=[(10,),
-              (20, 10),
-              (5, 8, 10)],
-        num_out=[20, ]
-    )
-    def test_Dense1(self, size, num_out):
-        f = brainstate.nn.Linear(10, num_out)
-        x = brainstate.random.random(size)
-        y = f(x)
-        self.assertTrue(y.shape == size[:-1] + (num_out,))
-
-
-class TestSparseMatrix(unittest.TestCase):
-    def test_csr(self):
-        data = brainstate.random.rand(10, 20)
-        data = data * (data > 0.9)
-        f = brainstate.nn.SparseLinear(u.sparse.CSR.fromdense(data))
-
-        x = brainstate.random.rand(10)
-        y = f(x)
-        self.assertTrue(
-            u.math.allclose(
-                y,
-                x @ data
-            )
-        )
-
-        x = brainstate.random.rand(5, 10)
-        y = f(x)
-        self.assertTrue(
-            u.math.allclose(
-                y,
-                x @ data
-            )
-        )
-
-    def test_csc(self):
-        data = brainstate.random.rand(10, 20)
-        data = data * (data > 0.9)
-        f = brainstate.nn.SparseLinear(u.sparse.CSC.fromdense(data))
-
-        x = brainstate.random.rand(10)
-        y = f(x)
-        self.assertTrue(
-            u.math.allclose(
-                y,
-                x @ data
-            )
-        )
-
-        x = brainstate.random.rand(5, 10)
-        y = f(x)
-        self.assertTrue(
-            u.math.allclose(
-                y,
-                x @ data
-            )
-        )
-
-    def test_coo(self):
-        data = brainstate.random.rand(10, 20)
-        data = data * (data > 0.9)
-        f = brainstate.nn.SparseLinear(u.sparse.COO.fromdense(data))
-
-        x = brainstate.random.rand(10)
-        y = f(x)
-        self.assertTrue(
-            u.math.allclose(
-                y,
-                x @ data
-            )
-        )
-
-        x = brainstate.random.rand(5, 10)
-        y = f(x)
-        self.assertTrue(
-            u.math.allclose(
-                y,
-                x @ data
-            )
-        )
+# 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 unittest
+
+import brainunit as u
+from absl.testing import parameterized
+
+import brainstate
+
+
+class TestDense(parameterized.TestCase):
+    @parameterized.product(
+        size=[(10,),
+              (20, 10),
+              (5, 8, 10)],
+        num_out=[20, ]
+    )
+    def test_Dense1(self, size, num_out):
+        f = brainstate.nn.Linear(10, num_out)
+        x = brainstate.random.random(size)
+        y = f(x)
+        self.assertTrue(y.shape == size[:-1] + (num_out,))
+
+
+class TestSparseMatrix(unittest.TestCase):
+    def test_csr(self):
+        data = brainstate.random.rand(10, 20)
+        data = data * (data > 0.9)
+        f = brainstate.nn.SparseLinear(u.sparse.CSR.fromdense(data))
+
+        x = brainstate.random.rand(10)
+        y = f(x)
+        self.assertTrue(
+            u.math.allclose(
+                y,
+                x @ data
+            )
+        )
+
+        x = brainstate.random.rand(5, 10)
+        y = f(x)
+        self.assertTrue(
+            u.math.allclose(
+                y,
+                x @ data
+            )
+        )
+
+    def test_csc(self):
+        data = brainstate.random.rand(10, 20)
+        data = data * (data > 0.9)
+        f = brainstate.nn.SparseLinear(u.sparse.CSC.fromdense(data))
+
+        x = brainstate.random.rand(10)
+        y = f(x)
+        self.assertTrue(
+            u.math.allclose(
+                y,
+                x @ data
+            )
+        )
+
+        x = brainstate.random.rand(5, 10)
+        y = f(x)
+        self.assertTrue(
+            u.math.allclose(
+                y,
+                x @ data
+            )
+        )
+
+    def test_coo(self):
+        data = brainstate.random.rand(10, 20)
+        data = data * (data > 0.9)
+        f = brainstate.nn.SparseLinear(u.sparse.COO.fromdense(data))
+
+        x = brainstate.random.rand(10)
+        y = f(x)
+        self.assertTrue(
+            u.math.allclose(
+                y,
+                x @ data
+            )
+        )
+
+        x = brainstate.random.rand(5, 10)
+        y = f(x)
+        self.assertTrue(
+            u.math.allclose(
+                y,
+                x @ data
+            )
+        )

brainstate/nn/_ltp.py

@@ -1,28 +1,28 @@
-# Copyright 2025 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 -*-
-
-
-from ._synapse import Synapse
-
-__all__ = [
-    'LongTermPlasticity',
-]
-
-
-class LongTermPlasticity(Synapse):
-    pass
-
+# Copyright 2025 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 -*-
+
+
+from ._synapse import Synapse
+
+__all__ = [
+    'LongTermPlasticity',
+]
+
+
+class LongTermPlasticity(Synapse):
+    pass
+