brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/nn/_interaction/_conv_test.py

@@ -0,0 +1,239 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import pytest
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import brainstate as bst
+
+
+class TestConv(parameterized.TestCase):
+    def test_Conv2D_img(self):
+        img = jnp.zeros((2, 200, 198, 4))
+        for k in range(4):
+            x = 30 + 60 * k
+            y = 20 + 60 * k
+            img = img.at[0, x:x + 10, y:y + 10, k].set(1.0)
+            img = img.at[1, x:x + 20, y:y + 20, k].set(3.0)
+
+        net = bst.nn.Conv2d((200, 198, 4), out_channels=32, kernel_size=(3, 3),
+                            stride=(2, 1), padding='VALID', groups=4)
+        out = net(img)
+        print("out shape: ", out.shape)
+        self.assertEqual(out.shape, (2, 99, 196, 32))
+        # print("First output channel:")
+        # plt.figure(figsize=(10, 10))
+        # plt.imshow(np.array(img)[0, :, :, 0])
+        # plt.show()
+
+    def test_conv1D(self):
+        model = bst.nn.Conv1d((5, 3), out_channels=32, kernel_size=(3,))
+        input = jnp.ones((2, 5, 3))
+        out = model(input)
+        print("out shape: ", out.shape)
+        self.assertEqual(out.shape, (2, 5, 32))
+        # print("First output channel:")
+        # plt.figure(figsize=(10, 10))
+        # plt.imshow(np.array(out)[0, :, :])
+        # plt.show()
+
+    def test_conv2D(self):
+        model = bst.nn.Conv2d((5, 5, 3), out_channels=32, kernel_size=(3, 3))
+        input = jnp.ones((2, 5, 5, 3))
+
+        out = model(input)
+        print("out shape: ", out.shape)
+        self.assertEqual(out.shape, (2, 5, 5, 32))
+
+    def test_conv3D(self):
+        model = bst.nn.Conv3d((5, 5, 5, 3), out_channels=32, kernel_size=(3, 3, 3))
+        input = jnp.ones((2, 5, 5, 5, 3))
+        out = model(input)
+        print("out shape: ", out.shape)
+        self.assertEqual(out.shape, (2, 5, 5, 5, 32))
+
+
+@pytest.mark.skip(reason="not implemented yet")
+class TestConvTranspose1d(parameterized.TestCase):
+    def test_conv_transpose(self):
+
+        x = jnp.ones((1, 8, 3))
+        for use_bias in [True, False]:
+            conv_transpose_module = bst.nn.ConvTranspose1d(
+                in_channels=3,
+                out_channels=4,
+                kernel_size=(3,),
+                padding='VALID',
+                w_initializer=bst.init.Constant(1.),
+                b_initializer=bst.init.Constant(1.) if use_bias else None,
+            )
+            self.assertEqual(conv_transpose_module.w.shape, (3, 3, 4))
+            y = conv_transpose_module(x)
+            print(y.shape)
+            correct_ans = jnp.array([[[4., 4., 4., 4.],
+                                      [7., 7., 7., 7.],
+                                      [10., 10., 10., 10.],
+                                      [10., 10., 10., 10.],
+                                      [10., 10., 10., 10.],
+                                      [10., 10., 10., 10.],
+                                      [10., 10., 10., 10.],
+                                      [10., 10., 10., 10.],
+                                      [7., 7., 7., 7.],
+                                      [4., 4., 4., 4.]]])
+            if not use_bias:
+                correct_ans -= 1.
+            self.assertTrue(jnp.allclose(y, correct_ans))
+
+    def test_single_input_masked_conv_transpose(self):
+
+        x = jnp.ones((1, 8, 3))
+        m = jnp.tril(jnp.ones((3, 3, 4)))
+        conv_transpose_module = bst.nn.ConvTranspose1d(
+            in_channels=3,
+            out_channels=4,
+            kernel_size=(3,),
+            padding='VALID',
+            mask=m,
+            w_initializer=bst.init.Constant(),
+            b_initializer=bst.init.Constant(),
+        )
+        self.assertEqual(conv_transpose_module.w.shape, (3, 3, 4))
+        y = conv_transpose_module(x)
+        print(y.shape)
+        correct_ans = jnp.array([[[4., 3., 2., 1.],
+                                  [7., 5., 3., 1.],
+                                  [10., 7., 4., 1.],
+                                  [10., 7., 4., 1.],
+                                  [10., 7., 4., 1.],
+                                  [10., 7., 4., 1.],
+                                  [10., 7., 4., 1.],
+                                  [10., 7., 4., 1.],
+                                  [7., 5., 3., 1.],
+                                  [4., 3., 2., 1.]]])
+        self.assertTrue(jnp.allclose(y, correct_ans))
+
+    def test_computation_padding_same(self):
+
+        data = jnp.ones([1, 3, 1])
+        for use_bias in [True, False]:
+            net = bst.nn.ConvTranspose1d(
+                in_channels=1,
+                out_channels=1,
+                kernel_size=3,
+                stride=1,
+                padding="SAME",
+                w_initializer=bst.init.Constant(),
+                b_initializer=bst.init.Constant() if use_bias else None,
+            )
+            out = net(data)
+            self.assertEqual(out.shape, (1, 3, 1))
+            out = jnp.squeeze(out, axis=(0, 2))
+            expected_out = jnp.asarray([2, 3, 2])
+            if use_bias:
+                expected_out += 1
+            self.assertTrue(jnp.allclose(out, expected_out, rtol=1e-5))
+
+
+@pytest.mark.skip(reason="not implemented yet")
+class TestConvTranspose2d(parameterized.TestCase):
+    def test_conv_transpose(self):
+
+        x = jnp.ones((1, 8, 8, 3))
+        for use_bias in [True, False]:
+            conv_transpose_module = bst.nn.ConvTranspose2d(
+                in_channels=3,
+                out_channels=4,
+                kernel_size=(3, 3),
+                padding='VALID',
+                w_initializer=bst.init.Constant(),
+                b_initializer=bst.init.Constant() if use_bias else None,
+            )
+        self.assertEqual(conv_transpose_module.w.shape, (3, 3, 3, 4))
+        y = conv_transpose_module(x)
+        print(y.shape)
+
+    def test_single_input_masked_conv_transpose(self):
+
+        x = jnp.ones((1, 8, 8, 3))
+        m = jnp.tril(jnp.ones((3, 3, 3, 4)))
+        conv_transpose_module = bst.nn.ConvTranspose2d(
+            in_channels=3,
+            out_channels=4,
+            kernel_size=(3, 3),
+            padding='VALID',
+            mask=m,
+            w_initializer=bst.init.Constant(),
+        )
+        y = conv_transpose_module(x)
+        print(y.shape)
+
+    def test_computation_padding_same(self):
+
+        x = jnp.ones((1, 8, 8, 3))
+        for use_bias in [True, False]:
+            conv_transpose_module = bst.nn.ConvTranspose2d(
+                in_channels=3,
+                out_channels=4,
+                kernel_size=(3, 3),
+                stride=1,
+                padding='SAME',
+                w_initializer=bst.init.Constant(),
+                b_initializer=bst.init.Constant() if use_bias else None,
+            )
+        y = conv_transpose_module(x)
+        print(y.shape)
+
+
+@pytest.mark.skip(reason="not implemented yet")
+class TestConvTranspose3d(parameterized.TestCase):
+    def test_conv_transpose(self):
+
+        x = jnp.ones((1, 8, 8, 8, 3))
+        for use_bias in [True, False]:
+            conv_transpose_module = bst.nn.ConvTranspose3d(
+                in_channels=3,
+                out_channels=4,
+                kernel_size=(3, 3, 3),
+                padding='VALID',
+                w_initializer=bst.init.Constant(),
+                b_initializer=bst.init.Constant() if use_bias else None,
+            )
+        y = conv_transpose_module(x)
+        print(y.shape)
+
+    def test_single_input_masked_conv_transpose(self):
+
+        x = jnp.ones((1, 8, 8, 8, 3))
+        m = jnp.tril(jnp.ones((3, 3, 3, 3, 4)))
+        conv_transpose_module = bst.nn.ConvTranspose3d(
+            in_channels=3,
+            out_channels=4,
+            kernel_size=(3, 3, 3),
+            padding='VALID',
+            mask=m,
+            w_initializer=bst.init.Constant(),
+        )
+        y = conv_transpose_module(x)
+        print(y.shape)
+
+    def test_computation_padding_same(self):
+
+        x = jnp.ones((1, 8, 8, 8, 3))
+        for use_bias in [True, False]:
+            conv_transpose_module = bst.nn.ConvTranspose3d(
+                in_channels=3,
+                out_channels=4,
+                kernel_size=(3, 3, 3),
+                stride=1,
+                padding='SAME',
+                w_initializer=bst.init.Constant(),
+                b_initializer=bst.init.Constant() if use_bias else None,
+            )
+        y = conv_transpose_module(x)
+        print(y.shape)
+
+if __name__ == '__main__':
+    absltest.main()

brainstate/nn/_interaction/_embedding.py

@@ -0,0 +1,59 @@
+# 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 __future__ import annotations
+
+from typing import Optional, Callable, Union
+
+from brainstate import init
+from brainstate._state import ParamState
+from brainstate.nn._module import Module
+from brainstate.typing import ArrayLike
+
+__all__ = [
+    'Embedding',
+]
+
+
+class Embedding(Module):
+    r"""
+    A simple lookup table that stores embeddings of a fixed size.
+
+    Args:
+      num_embeddings: Size of embedding dictionary. Must be non-negative.
+      embedding_size: Size of each embedding vector. Must be non-negative.
+      embedding_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
+    """
+
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_size: int,
+        embedding_init: Union[Callable, ArrayLike] = init.LecunUniform(),
+        name: Optional[str] = None,
+    ):
+        super().__init__(name=name)
+        if num_embeddings < 0:
+            raise ValueError("num_embeddings must not be negative.")
+        if embedding_size < 0:
+            raise ValueError("embedding_size must not be negative.")
+        self.num_embeddings = num_embeddings
+        self.embedding_size = embedding_size
+        self.out_size = (embedding_size,)
+
+        weight = init.param(embedding_init, (self.num_embeddings, self.embedding_size))
+        self.weight = ParamState(weight)
+
+    def update(self, indices: ArrayLike):
+        return self.weight.value[indices]