brainstate 0.0.1__py2.py3-none-any.whl

brainstate/nn/_poolings_test.py

@@ -0,0 +1,231 @@
+# -*- coding: utf-8 -*-
+
+
+import jax
+import numpy as np
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import brainstate as bst
+import brainstate.nn as nn
+
+
+class TestFlatten(parameterized.TestCase):
+  def test_flatten1(self):
+    for size in [
+      (16, 32, 32, 8),
+      (32, 8),
+      (10, 20, 30),
+    ]:
+      arr = bst.random.rand(*size)
+      f = nn.Flatten(start_dim=0)
+      out = f(arr)
+      self.assertTrue(out.shape == (np.prod(size),))
+
+  def test_flatten2(self):
+    for size in [
+      (16, 32, 32, 8),
+      (32, 8),
+      (10, 20, 30),
+    ]:
+      arr = bst.random.rand(*size)
+      f = nn.Flatten(start_dim=1)
+      out = f(arr)
+      self.assertTrue(out.shape == (size[0], np.prod(size[1:])))
+
+  def test_flatten3(self):
+    size = (16, 32, 32, 8)
+    arr = bst.random.rand(*size)
+    f = nn.Flatten(start_dim=0, in_size=(32, 8))
+    out = f(arr)
+    self.assertTrue(out.shape == (16, 32, 32 * 8))
+
+  def test_flatten4(self):
+    size = (16, 32, 32, 8)
+    arr = bst.random.rand(*size)
+    f = nn.Flatten(start_dim=1, in_size=(32, 32, 8))
+    out = f(arr)
+    self.assertTrue(out.shape == (16, 32, 32 * 8))
+
+
+class TestUnflatten(parameterized.TestCase):
+  pass
+
+
+class TestPool(parameterized.TestCase):
+  def __init__(self, *args, **kwargs):
+    super().__init__(*args, **kwargs)
+
+  def test_MaxPool2d_v1(self):
+    arr = bst.random.rand(16, 32, 32, 8)
+
+    out = nn.MaxPool2d(2, 2, channel_axis=-1)(arr)
+    self.assertTrue(out.shape == (16, 16, 16, 8))
+
+    out = nn.MaxPool2d(2, 2, channel_axis=None)(arr)
+    self.assertTrue(out.shape == (16, 32, 16, 4))
+
+    out = nn.MaxPool2d(2, 2, channel_axis=None, padding=1)(arr)
+    self.assertTrue(out.shape == (16, 32, 17, 5))
+
+    out = nn.MaxPool2d(2, 2, channel_axis=None, padding=(2, 1))(arr)
+    self.assertTrue(out.shape == (16, 32, 18, 5))
+
+    out = nn.MaxPool2d(2, 2, channel_axis=-1, padding=(1, 1))(arr)
+    self.assertTrue(out.shape == (16, 17, 17, 8))
+
+    out = nn.MaxPool2d(2, 2, channel_axis=2, padding=(1, 1))(arr)
+    self.assertTrue(out.shape == (16, 17, 32, 5))
+    bst.util.clear_buffer_memory()
+
+  def test_AvgPool2d_v1(self):
+    arr = bst.random.rand(16, 32, 32, 8)
+
+    out = nn.AvgPool2d(2, 2, channel_axis=-1)(arr)
+    self.assertTrue(out.shape == (16, 16, 16, 8))
+
+    out = nn.AvgPool2d(2, 2, channel_axis=None)(arr)
+    self.assertTrue(out.shape == (16, 32, 16, 4))
+
+    out = nn.AvgPool2d(2, 2, channel_axis=None, padding=1)(arr)
+    self.assertTrue(out.shape == (16, 32, 17, 5))
+
+    out = nn.AvgPool2d(2, 2, channel_axis=None, padding=(2, 1))(arr)
+    self.assertTrue(out.shape == (16, 32, 18, 5))
+
+    out = nn.AvgPool2d(2, 2, channel_axis=-1, padding=(1, 1))(arr)
+    self.assertTrue(out.shape == (16, 17, 17, 8))
+
+    out = nn.AvgPool2d(2, 2, channel_axis=2, padding=(1, 1))(arr)
+    self.assertTrue(out.shape == (16, 17, 32, 5))
+    bst.util.clear_buffer_memory()
+
+  @parameterized.named_parameters(
+    dict(testcase_name=f'target_size={target_size}',
+         target_size=target_size)
+    for target_size in [10, 9, 8, 7, 6]
+  )
+  def test_adaptive_pool1d(self, target_size):
+    from brainstate.nn._poolings import _adaptive_pool1d
+
+    arr = bst.random.rand(100)
+    op = jax.numpy.mean
+
+    out = _adaptive_pool1d(arr, target_size, op)
+    print(out.shape)
+    self.assertTrue(out.shape == (target_size,))
+
+    out = _adaptive_pool1d(arr, target_size, op)
+    print(out.shape)
+    self.assertTrue(out.shape == (target_size,))
+    bst.util.clear_buffer_memory()
+
+  def test_AdaptiveAvgPool2d_v1(self):
+    input = bst.random.randn(64, 8, 9)
+
+    output = nn.AdaptiveAvgPool2d((5, 7), channel_axis=0)(input)
+    self.assertTrue(output.shape == (64, 5, 7))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=0)(input)
+    self.assertTrue(output.shape == (64, 2, 3))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=-1)(input)
+    self.assertTrue(output.shape == (2, 3, 9))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=1)(input)
+    self.assertTrue(output.shape == (2, 8, 3))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=None)(input)
+    self.assertTrue(output.shape == (64, 2, 3))
+    bst.util.clear_buffer_memory()
+
+  def test_AdaptiveAvgPool2d_v2(self):
+    bst.random.seed()
+    input = bst.random.randn(128, 64, 32, 16)
+
+    output = nn.AdaptiveAvgPool2d((5, 7), channel_axis=0)(input)
+    self.assertTrue(output.shape == (128, 64, 5, 7))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=0)(input)
+    self.assertTrue(output.shape == (128, 64, 2, 3))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=-1)(input)
+    self.assertTrue(output.shape == (128, 2, 3, 16))
+
+    output = nn.AdaptiveAvgPool2d((2, 3), channel_axis=1)(input)
+    self.assertTrue(output.shape == (128, 64, 2, 3))
+    print()
+    bst.util.clear_buffer_memory()
+
+  def test_AdaptiveAvgPool3d_v1(self):
+    input = bst.random.randn(10, 128, 64, 32)
+    net = nn.AdaptiveAvgPool3d(target_size=[6, 5, 3], channel_axis=0)
+    output = net(input)
+    self.assertTrue(output.shape == (10, 6, 5, 3))
+    bst.util.clear_buffer_memory()
+
+  def test_AdaptiveAvgPool3d_v2(self):
+    input = bst.random.randn(10, 20, 128, 64, 32)
+    net = nn.AdaptiveAvgPool3d(target_size=[6, 5, 3])
+    output = net(input)
+    self.assertTrue(output.shape == (10, 6, 5, 3, 32))
+    bst.util.clear_buffer_memory()
+
+  @parameterized.product(
+    axis=(-1, 0, 1)
+  )
+  def test_AdaptiveMaxPool1d_v1(self, axis):
+    input = bst.random.randn(32, 16)
+    net = nn.AdaptiveMaxPool1d(target_size=4, channel_axis=axis)
+    output = net(input)
+    bst.util.clear_buffer_memory()
+
+  @parameterized.product(
+    axis=(-1, 0, 1, 2)
+  )
+  def test_AdaptiveMaxPool1d_v2(self, axis):
+    input = bst.random.randn(2, 32, 16)
+    net = nn.AdaptiveMaxPool1d(target_size=4, channel_axis=axis)
+    output = net(input)
+    bst.util.clear_buffer_memory()
+
+  @parameterized.product(
+    axis=(-1, 0, 1, 2)
+  )
+  def test_AdaptiveMaxPool2d_v1(self, axis):
+    input = bst.random.randn(32, 16, 12)
+    net = nn.AdaptiveAvgPool2d(target_size=[5, 4], channel_axis=axis)
+    output = net(input)
+    bst.util.clear_buffer_memory()
+
+  @parameterized.product(
+    axis=(-1, 0, 1, 2, 3)
+  )
+  def test_AdaptiveMaxPool2d_v2(self, axis):
+    input = bst.random.randn(2, 32, 16, 12)
+    net = nn.AdaptiveAvgPool2d(target_size=[5, 4], channel_axis=axis)
+    output = net(input)
+    bst.util.clear_buffer_memory()
+
+  @parameterized.product(
+    axis=(-1, 0, 1, 2, 3)
+  )
+  def test_AdaptiveMaxPool3d_v1(self, axis):
+    input = bst.random.randn(2, 128, 64, 32)
+    net = nn.AdaptiveMaxPool3d(target_size=[6, 5, 4], channel_axis=axis)
+    output = net(input)
+    print()
+    bst.util.clear_buffer_memory()
+
+  @parameterized.product(
+    axis=(-1, 0, 1, 2, 3, 4)
+  )
+  def test_AdaptiveMaxPool3d_v1(self, axis):
+    input = bst.random.randn(2, 128, 64, 32, 16)
+    net = nn.AdaptiveMaxPool3d(target_size=[6, 5, 4], channel_axis=axis)
+    output = net(input)
+    bst.util.clear_buffer_memory()
+
+
+if __name__ == '__main__':
+  absltest.main()

brainstate/nn/_projection/__init__.py

@@ -0,0 +1,32 @@
+# 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.
+# ==============================================================================
+
+"""
+
+This module defines the basic classes for synaptic projections.
+
+"""
+
+from ._align_post import *
+from ._align_post import __all__ as align_post_all
+from ._align_pre import *
+from ._align_pre import __all__ as align_pre_all
+from ._delta import *
+from ._delta import __all__ as delta_all
+from ._vanilla import *
+from ._vanilla import __all__ as vanilla_all
+
+__all__ = align_post_all + align_pre_all + delta_all + vanilla_all
+del align_post_all, align_pre_all, delta_all, vanilla_all