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

brainstate/nn/_interaction/_poolings_test.py

@@ -0,0 +1,219 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+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_axis=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_axis=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_axis=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_axis=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))
+
+    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))
+
+    @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._interaction._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,))
+
+    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))
+
+    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()
+
+    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))
+
+    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))
+
+    @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)
+
+    @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)
+
+    @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)
+
+    @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)
+
+    @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()
+
+    @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)
+
+
+if __name__ == '__main__':
+    absltest.main()

brainstate/nn/_module.py

@@ -0,0 +1,328 @@
+# 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 -*-
+
+
+"""
+All the basic classes for neural networks in ``brainstate``.
+
+The basic classes include:
+
+- ``Module``: The base class for all the objects in the ecosystem.
+- ``Sequential``: The class for a sequential of modules, which update the modules sequentially.
+
+"""
+from __future__ import annotations
+
+import warnings
+from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING
+
+import numpy as np
+
+from brainstate._state import State
+from brainstate.graph import Node, states, nodes, flatten
+from brainstate.mixin import ParamDescriber, ParamDesc
+from brainstate.typing import PathParts
+from brainstate.util import FlattedDict, NestedDict
+
+# maximum integer
+max_int = np.iinfo(np.int32).max
+
+__all__ = [
+    'Module', 'ElementWiseBlock', 'Sequential',
+]
+
+
+class Module(Node, ParamDesc):
+    """
+    The Module class for the whole ecosystem.
+
+    The ``Module`` is the base class for all the objects in the ecosystem. It
+    provides the basic functionalities for the objects, including:
+
+    - ``states()``: Collect all states in this node and the children nodes.
+    - ``nodes()``: Collect all children nodes.
+    - ``update()``: The function to specify the updating rule.
+    - ``init_state()``: State initialization function.
+    - ``reset_state()``: State resetting function.
+
+    """
+
+    __module__ = 'brainstate.nn'
+
+    _in_size: Optional[Tuple[int, ...]]
+    _out_size: Optional[Tuple[int, ...]]
+    _name: Optional[str]
+
+    if not TYPE_CHECKING:
+        def __init__(self, name: str = None):
+            # check the name
+            if name is not None:
+                assert isinstance(name, str), f'The name must be a string, but we got {type(name)}: {name}'
+            self._name = name
+
+            # input and output size
+            self._in_size = None
+            self._out_size = None
+
+    @property
+    def name(self):
+        """Name of the model."""
+        return self._name
+
+    @name.setter
+    def name(self, name: str = None):
+        raise AttributeError('The name of the model is read-only.')
+
+    @property
+    def in_size(self) -> Tuple[int, ...]:
+        return self._in_size
+
+    @in_size.setter
+    def in_size(self, in_size: Sequence[int] | int):
+        if isinstance(in_size, int):
+            in_size = (in_size,)
+        assert isinstance(in_size, (tuple, list)), f"Invalid type of in_size: {type(in_size)}"
+        self._in_size = tuple(in_size)
+
+    @property
+    def out_size(self) -> Tuple[int, ...]:
+        return self._out_size
+
+    @out_size.setter
+    def out_size(self, out_size: Sequence[int] | int):
+        if isinstance(out_size, int):
+            out_size = (out_size,)
+        assert isinstance(out_size, (tuple, list)), f"Invalid type of out_size: {type(out_size)}"
+        self._out_size = tuple(out_size)
+
+    def update(self, *args, **kwargs):
+        """
+        The function to specify the updating rule.
+        """
+        raise NotImplementedError(f'Subclass of {self.__class__.__name__} must implement "update" function.')
+
+    def __call__(self, *args, **kwargs):
+        return self.update(*args, **kwargs)
+
+    def __rrshift__(self, other):
+        """
+        Support using right shift operator to call modules.
+
+        Examples
+        --------
+
+        >>> import brainstate as bst
+        >>> x = bst.random.rand((10, 10))
+        >>> l = bst.nn.Dropout(0.5)
+        >>> y = x >> l
+        """
+        return self.__call__(other)
+
+    def states(
+        self,
+        *filters,
+        allowed_hierarchy: Tuple[int, int] = (0, max_int),
+        level: int = None,
+    ) -> FlattedDict[PathParts, State] | Tuple[FlattedDict[PathParts, State], ...]:
+        """
+        Collect all states in this node and the children nodes.
+
+        Parameters
+        ----------
+        filters : Any
+          The filters to select the states.
+        allowed_hierarchy : tuple of int
+          The hierarchy of the states to be collected.
+        level : int
+          The level of the states to be collected. Has been deprecated.
+
+        Returns
+        -------
+        states : FlattedDict, tuple of FlattedDict
+          The collection contained (the path, the state).
+        """
+        if level is not None:
+            allowed_hierarchy = (0, level)
+            warnings.warn('The "level" argument is deprecated. Please use "allowed_hierarchy" instead.',
+                          DeprecationWarning)
+
+        return states(self, *filters, allowed_hierarchy=allowed_hierarchy)
+
+    def state_trees(
+        self,
+        *filters,
+    ) -> NestedDict[PathParts, State] | Tuple[NestedDict[PathParts, State], ...]:
+        """
+        Collect all states in this node and the children nodes.
+
+        Parameters
+        ----------
+        filters : tuple
+          The filters to select the states.
+
+        Returns
+        -------
+        states : FlattedDict, tuple of FlattedDict
+          The collection contained (the path, the state).
+        """
+        graph_def, state_tree = flatten(self)
+        if len(filters):
+            return state_tree.filter(*filters)
+        return state_tree
+
+    def nodes(
+        self,
+        *filters,
+        allowed_hierarchy: Tuple[int, int] = (0, max_int),
+        level: int = None,
+    ) -> FlattedDict[PathParts, Node] | Tuple[FlattedDict[PathParts, Node], ...]:
+        """
+        Collect all children nodes.
+
+        Parameters
+        ----------
+        filters : Any
+          The filters to select the states.
+        allowed_hierarchy : tuple of int
+          The hierarchy of the states to be collected.
+        level : int
+          The level of the states to be collected. Has been deprecated.
+
+        Returns
+        -------
+        nodes : FlattedDict, tuple of FlattedDict
+          The collection contained (the path, the node).
+        """
+        if level is not None:
+            allowed_hierarchy = (0, level)
+            warnings.warn('The "level" argument is deprecated. Please use "allowed_hierarchy" instead.',
+                          DeprecationWarning)
+
+        return nodes(self, *filters, allowed_hierarchy=allowed_hierarchy)
+
+    def init_state(self, *args, **kwargs):
+        """
+        State initialization function.
+        """
+        pass
+
+    def reset_state(self, *args, **kwargs):
+        """
+        State resetting function.
+        """
+        pass
+
+    def __leaf_fn__(self, name, value):
+        if name in ['_in_size', '_out_size', '_name']:
+            return (name, value) if value is None else (name[1:], value)  # skip the first `_`
+        return name, value
+
+
+class ElementWiseBlock(Module):
+    __module__ = 'brainstate.nn'
+
+
+class Sequential(Module):
+    """
+    A sequential `input-output` module.
+
+    Modules will be added to it in the order they are passed in the
+    constructor. Alternatively, an ``dict`` of modules can be
+    passed in. The ``update()`` method of ``Sequential`` accepts any
+    input and forwards it to the first module it contains. It then
+    "chains" outputs to inputs sequentially for each subsequent module,
+    finally returning the output of the last module.
+
+    The value a ``Sequential`` provides over manually calling a sequence
+    of modules is that it allows treating the whole container as a
+    single module, such that performing a transformation on the
+    ``Sequential`` applies to each of the modules it stores (which are
+    each a registered submodule of the ``Sequential``).
+
+    What's the difference between a ``Sequential`` and a
+    :py:class:`Container`? A ``Container`` is exactly what it
+    sounds like--a container to store :py:class:`DynamicalSystem` s!
+    On the other hand, the layers in a ``Sequential`` are connected
+    in a cascading way.
+
+    Examples
+    --------
+
+    >>> import jax
+    >>> import brainstate as bst
+    >>> import brainstate.nn as nn
+    >>>
+    >>> # composing ANN models
+    >>> l = nn.Sequential(nn.Linear(100, 10),
+    >>>                   jax.nn.relu,
+    >>>                   nn.Linear(10, 2))
+    >>> l(bst.random.random((256, 100)))
+
+    Args:
+      modules_as_tuple: The children modules.
+      modules_as_dict: The children modules.
+      name: The object name.
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(self, first: Module, *layers):
+        super().__init__()
+        self.layers = []
+
+        # add all modules
+        assert isinstance(first, Module), 'The first module should be an instance of Module.'
+        in_size = first.out_size
+        self.layers.append(first)
+        for module in layers:
+            module, in_size = _format_module(module, in_size)
+            self.layers.append(module)
+
+        # the input and output shape
+        if first.in_size is not None:
+            self.in_size = first.in_size
+        self.out_size = tuple(in_size)
+
+    def update(self, x):
+        """Update function of a sequential model.
+        """
+        for m in self.layers:
+            x = m(x)
+        return x
+
+    def __getitem__(self, key: Union[int, slice]):
+        if isinstance(key, slice):
+            return Sequential(*self.layers[key])
+        elif isinstance(key, int):
+            return self.layers[key]
+        elif isinstance(key, (tuple, list)):
+            return Sequential(*[self.layers[k] for k in key])
+        else:
+            raise KeyError(f'Unknown type of key: {type(key)}')
+
+
+def _format_module(module, in_size):
+    if isinstance(module, ParamDescriber):
+        module = module(in_size=in_size)
+        assert isinstance(module, Module), 'The module should be an instance of Module.'
+        out_size = module.out_size
+    elif isinstance(module, ElementWiseBlock):
+        out_size = in_size
+    elif isinstance(module, Module):
+        out_size = module.out_size
+    else:
+        raise TypeError(f"Unsupported type {type(module)}. ")
+    return module, out_size