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

brainstate/nn/_poolings.py

@@ -1,1229 +0,0 @@
-# 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 -*-
-
-from __future__ import annotations
-
-import functools
-from typing import Sequence, Optional
-from typing import Union, Tuple, Callable, List
-
-import brainunit as u
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from ._base import DnnLayer, ExplicitInOutSize
-from .. import environ
-from ..mixin import Mode
-from brainstate.typing import Size
-
-__all__ = [
-  'Flatten', 'Unflatten',
-
-  'AvgPool1d', 'AvgPool2d', 'AvgPool3d',
-  'MaxPool1d', 'MaxPool2d', 'MaxPool3d',
-
-  'AdaptiveAvgPool1d', 'AdaptiveAvgPool2d', 'AdaptiveAvgPool3d',
-  'AdaptiveMaxPool1d', 'AdaptiveMaxPool2d', 'AdaptiveMaxPool3d',
-]
-
-
-class Flatten(DnnLayer, ExplicitInOutSize):
-  r"""
-  Flattens a contiguous range of dims into a tensor. For use with :class:`~nn.Sequential`.
-
-  Shape:
-      - Input: :math:`(*, S_{\text{start}},..., S_{i}, ..., S_{\text{end}}, *)`,'
-        where :math:`S_{i}` is the size at dimension :math:`i` and :math:`*` means any
-        number of dimensions including none.
-      - Output: :math:`(*, \prod_{i=\text{start}}^{\text{end}} S_{i}, *)`.
-
-  Args:
-      in_size: Sequence of int. The shape of the input tensor.
-      start_axis: first dim to flatten (default = 1).
-      end_axis: last dim to flatten (default = -1).
-
-  Examples::
-      >>> import brainstate as bst
-      >>> inp = bst.random.randn(32, 1, 5, 5)
-      >>> # With default parameters
-      >>> m = Flatten()
-      >>> output = m(inp)
-      >>> output.shape
-      (32, 25)
-      >>> # With non-default parameters
-      >>> m = Flatten(0, 2)
-      >>> output = m(inp)
-      >>> output.shape
-      (160, 5)
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      start_axis: int = 0,
-      end_axis: int = -1,
-      in_size: Optional[Size] = None
-  ) -> None:
-    super().__init__()
-    self.start_axis = start_axis
-    self.end_axis = end_axis
-
-    if in_size is not None:
-      self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(u.math.flatten, start_axis=start_axis, end_axis=end_axis),
-                         jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
-      self.out_size = y.shape
-
-  def update(self, x):
-    if self._in_size is None:
-      start_axis = self.start_axis if self.start_axis >= 0 else x.ndim + self.start_axis
-    else:
-      assert x.ndim >= len(self.in_size), 'Input tensor has fewer dimensions than the expected shape.'
-      dim_diff = x.ndim - len(self.in_size)
-      if self.in_size != x.shape[dim_diff:]:
-        raise ValueError(f'Input tensor has shape {x.shape}, but expected shape {self.in_size}.')
-      if self.start_axis >= 0:
-        start_axis = self.start_axis + dim_diff
-      else:
-        start_axis = x.ndim + self.start_axis
-    return u.math.flatten(x, start_axis, self.end_axis)
-
-  def __repr__(self) -> str:
-    return f'{self.__class__.__name__}(start_axis={self.start_axis}, end_axis={self.end_axis})'
-
-
-class Unflatten(DnnLayer, ExplicitInOutSize):
-  r"""
-  Unflatten a tensor dim expanding it to a desired shape. For use with :class:`~nn.Sequential`.
-
-  * :attr:`dim` specifies the dimension of the input tensor to be unflattened, and it can
-    be either `int` or `str` when `Tensor` or `NamedTensor` is used, respectively.
-
-  * :attr:`unflattened_size` is the new shape of the unflattened dimension of the tensor and it can be
-    a `tuple` of ints or a `list` of ints or `torch.Size` for `Tensor` input;  a `NamedShape`
-    (tuple of `(name, size)` tuples) for `NamedTensor` input.
-
-  Shape:
-      - Input: :math:`(*, S_{\text{dim}}, *)`, where :math:`S_{\text{dim}}` is the size at
-        dimension :attr:`dim` and :math:`*` means any number of dimensions including none.
-      - Output: :math:`(*, U_1, ..., U_n, *)`, where :math:`U` = :attr:`unflattened_size` and
-        :math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
-
-  Args:
-      axis: int, Dimension to be unflattened.
-      sizes: Sequence of int. New shape of the unflattened dimension.
-      in_size: Sequence of int. The shape of the input tensor.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      axis: int,
-      sizes: Size,
-      mode: Mode = None,
-      name: str = None,
-      in_size: Optional[Size] = None
-  ) -> None:
-    super().__init__(mode=mode, name=name)
-
-    self.axis = axis
-    self.sizes = sizes
-    if isinstance(sizes, (tuple, list)):
-      for idx, elem in enumerate(sizes):
-        if not isinstance(elem, int):
-          raise TypeError("unflattened sizes must be tuple of ints, " +
-                          "but found element of type {} at pos {}".format(type(elem).__name__, idx))
-    else:
-      raise TypeError("unflattened sizes must be tuple or list, but found type {}".format(type(sizes).__name__))
-
-    if in_size is not None:
-      self.in_size = tuple(in_size)
-      y = jax.eval_shape(functools.partial(u.math.unflatten, axis=axis, sizes=sizes),
-                         jax.ShapeDtypeStruct(self.in_size, environ.dftype()))
-      self.out_size = y.shape
-
-  def update(self, x):
-    return u.math.unflatten(x, self.axis, self.sizes)
-
-  def __repr__(self):
-    return f'{self.__class__.__name__}(axis={self.axis}, sizes={self.sizes})'
-
-
-class _MaxPool(DnnLayer, ExplicitInOutSize):
-  def __init__(
-      self,
-      init_value: float,
-      computation: Callable,
-      pool_dim: int,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = None,
-      padding: Union[str, int, Tuple[int, ...], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(name=name, mode=mode)
-
-    self.init_value = init_value
-    self.computation = computation
-    self.pool_dim = pool_dim
-
-    # kernel_size
-    if isinstance(kernel_size, int):
-      kernel_size = (kernel_size,) * pool_dim
-    elif isinstance(kernel_size, Sequence):
-      assert isinstance(kernel_size, (tuple, list)), f'kernel_size should be a tuple, but got {type(kernel_size)}'
-      assert all([isinstance(x, int) for x in kernel_size]), f'kernel_size should be a tuple of ints. {kernel_size}'
-      if len(kernel_size) != pool_dim:
-        raise ValueError(f'kernel_size should a tuple with {pool_dim} ints, but got {len(kernel_size)}')
-    else:
-      raise TypeError(f'kernel_size should be a int or a tuple with {pool_dim} ints.')
-    self.kernel_size = kernel_size
-
-    # stride
-    if stride is None:
-      stride = kernel_size
-    if isinstance(stride, int):
-      stride = (stride,) * pool_dim
-    elif isinstance(stride, Sequence):
-      assert isinstance(stride, (tuple, list)), f'stride should be a tuple, but got {type(stride)}'
-      assert all([isinstance(x, int) for x in stride]), f'stride should be a tuple of ints. {stride}'
-      if len(stride) != pool_dim:
-        raise ValueError(f'stride should a tuple with {pool_dim} ints, but got {len(kernel_size)}')
-    else:
-      raise TypeError(f'stride should be a int or a tuple with {pool_dim} ints.')
-    self.stride = stride
-
-    # padding
-    if isinstance(padding, str):
-      if padding not in ("SAME", "VALID"):
-        raise ValueError(f"Invalid padding '{padding}', must be 'SAME' or 'VALID'.")
-    elif isinstance(padding, int):
-      padding = [(padding, padding) for _ in range(pool_dim)]
-    elif isinstance(padding, (list, tuple)):
-      if isinstance(padding[0], int):
-        if len(padding) == pool_dim:
-          padding = [(x, x) for x in padding]
-        else:
-          raise ValueError(f'If padding is a sequence of ints, it '
-                           f'should has the length of {pool_dim}.')
-      else:
-        if not all([isinstance(x, (tuple, list)) for x in padding]):
-          raise ValueError(f'padding should be sequence of Tuple[int, int]. {padding}')
-        if not all([len(x) == 2 for x in padding]):
-          raise ValueError(f"Each entry in padding must be tuple of 2 ints. {padding} ")
-        if len(padding) == 1:
-          padding = tuple(padding) * pool_dim
-        assert len(padding) == pool_dim, f'padding should has the length of {pool_dim}. {padding}'
-    else:
-      raise ValueError
-    self.padding = padding
-
-    # channel_axis
-    assert channel_axis is None or isinstance(channel_axis, int), \
-      f'channel_axis should be an int, but got {channel_axis}'
-    self.channel_axis = channel_axis
-
-    # in & out shapes
-    if in_size is not None:
-      in_size = tuple(in_size)
-      self.in_size = in_size
-      y = jax.eval_shape(self.update, jax.ShapeDtypeStruct((128,) + in_size, environ.dftype()))
-      self.out_size = y.shape[1:]
-
-  def update(self, x):
-    x_dim = self.pool_dim + (0 if self.channel_axis is None else 1)
-    if x.ndim < x_dim:
-      raise ValueError(f'Excepted input with >= {x_dim} dimensions, but got {x.ndim}.')
-    window_shape = self._infer_shape(x.ndim, self.kernel_size, 1)
-    stride = self._infer_shape(x.ndim, self.stride, 1)
-    padding = (self.padding if isinstance(self.padding, str) else
-               self._infer_shape(x.ndim, self.padding, element=(0, 0)))
-    r = jax.lax.reduce_window(
-      x,
-      init_value=self.init_value,
-      computation=self.computation,
-      window_dimensions=window_shape,
-      window_strides=stride,
-      padding=padding
-    )
-    return r
-
-  def _infer_shape(self, x_dim, inputs, element):
-    channel_axis = self.channel_axis
-    if channel_axis and not 0 <= abs(channel_axis) < x_dim:
-      raise ValueError(f"Invalid channel axis {channel_axis} for input with {x_dim} dimensions")
-    if channel_axis and channel_axis < 0:
-      channel_axis = x_dim + channel_axis
-    all_dims = list(range(x_dim))
-    if channel_axis is not None:
-      all_dims.pop(channel_axis)
-    pool_dims = all_dims[-self.pool_dim:]
-    results = [element] * x_dim
-    for i, dim in enumerate(pool_dims):
-      results[dim] = inputs[i]
-    return results
-
-
-class _AvgPool(_MaxPool):
-  def update(self, x):
-    x_dim = self.pool_dim + (0 if self.channel_axis is None else 1)
-    if x.ndim < x_dim:
-      raise ValueError(f'Excepted input with >= {x_dim} dimensions, but got {x.ndim}.')
-    dims = self._infer_shape(x.ndim, self.kernel_size, 1)
-    stride = self._infer_shape(x.ndim, self.stride, 1)
-    padding = (self.padding if isinstance(self.padding, str) else
-               self._infer_shape(x.ndim, self.padding, element=(0, 0)))
-    pooled = jax.lax.reduce_window(x,
-                                   init_value=self.init_value,
-                                   computation=self.computation,
-                                   window_dimensions=dims,
-                                   window_strides=stride,
-                                   padding=padding)
-    if padding == "VALID":
-      # Avoid the extra reduce_window.
-      return pooled / np.prod(dims)
-    else:
-      # Count the number of valid entries at each input point, then use that for
-      # computing average. Assumes that any two arrays of same shape will be
-      # padded the same.
-      window_counts = jax.lax.reduce_window(jnp.ones_like(x),
-                                            init_value=self.init_value,
-                                            computation=self.computation,
-                                            window_dimensions=dims,
-                                            window_strides=stride,
-                                            padding=padding)
-      assert pooled.shape == window_counts.shape
-      return pooled / window_counts
-
-
-class MaxPool1d(_MaxPool):
-  r"""Applies a 1D max pooling over an input signal composed of several input planes.
-
-  In the simplest case, the output value of the layer with input size :math:`(N, L, C)`
-  and output :math:`(N, L_{out}, C)` can be precisely described as:
-
-  .. math::
-      out(N_i, k, C_j) = \max_{m=0, \ldots, \text{kernel\_size} - 1}
-              input(N_i, stride \times k + m, C_j)
-
-  If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides
-  for :attr:`padding` number of points. :attr:`dilation` is the stride between the elements within the
-  sliding window. This `link`_ has a nice visualization of the pooling parameters.
-
-  Shape:
-      - Input: :math:`(N, L_{in}, C)` or :math:`(L_{in}, C)`.
-      - Output: :math:`(N, L_{out}, C)` or :math:`(L_{out}, C)`, where
-
-        .. math::
-            L_{out} = \left\lfloor \frac{L_{in} + 2 \times \text{padding} - \text{dilation}
-                  \times (\text{kernel\_size} - 1) - 1}{\text{stride}} + 1\right\rfloor
-
-
-  Examples::
-
-      >>> import brainstate as bst
-      >>> # pool of size=3, stride=2
-      >>> m = MaxPool1d(3, stride=2, channel_axis=-1)
-      >>> input = bst.random.randn(20, 50, 16)
-      >>> output = m(input)
-      >>> output.shape
-      (20, 24, 16)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  kernel_size: int, sequence of int
-    An integer, or a sequence of integers defining the window to reduce over.
-  stride: int, sequence of int
-    An integer, or a sequence of integers, representing the inter-window stride (default: `(1, ..., 1)`).
-  padding: str, int, sequence of tuple
-    Either the string `'SAME'`, the string `'VALID'`, or a sequence
-    of n `(low, high)` integer pairs that give the padding to apply before
-    and after each spatial dimension.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  mode: Mode
-    The computation mode.
-  name: optional, str
-    The object name.
-
-  .. _link:
-        https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = None,
-      padding: Union[str, int, Tuple[int, ...], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(in_size=in_size,
-                     init_value=-jax.numpy.inf,
-                     computation=jax.lax.max,
-                     pool_dim=1,
-                     kernel_size=kernel_size,
-                     stride=stride,
-                     padding=padding,
-                     channel_axis=channel_axis,
-                     name=name,
-                     mode=mode)
-
-
-class MaxPool2d(_MaxPool):
-  r"""Applies a 2D max pooling over an input signal composed of several input planes.
-
-  In the simplest case, the output value of the layer with input size :math:`(N, H, W, C)`,
-  output :math:`(N, H_{out}, W_{out}, C)` and :attr:`kernel_size` :math:`(kH, kW)`
-  can be precisely described as:
-
-  .. math::
-      \begin{aligned}
-          out(N_i, h, w, C_j) ={} & \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \\
-                                  & \text{input}(N_i, \text{stride[0]} \times h + m,
-                                                 \text{stride[1]} \times w + n, C_j)
-      \end{aligned}
-
-  If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides
-  for :attr:`padding` number of points. :attr:`dilation` controls the spacing between the kernel points.
-  It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
-
-
-  Shape:
-      - Input: :math:`(N, H_{in}, W_{in}, C)` or :math:`(H_{in}, W_{in}, C)`
-      - Output: :math:`(N, H_{out}, W_{out}, C)` or :math:`(H_{out}, W_{out}, C)`, where
-
-        .. math::
-            H_{out} = \left\lfloor\frac{H_{in} + 2 * \text{padding[0]} - \text{dilation[0]}
-                  \times (\text{kernel\_size[0]} - 1) - 1}{\text{stride[0]}} + 1\right\rfloor
-
-        .. math::
-            W_{out} = \left\lfloor\frac{W_{in} + 2 * \text{padding[1]} - \text{dilation[1]}
-                  \times (\text{kernel\_size[1]} - 1) - 1}{\text{stride[1]}} + 1\right\rfloor
-
-  Examples::
-
-      >>> import brainstate as bst
-      >>> # pool of square window of size=3, stride=2
-      >>> m = MaxPool2d(3, stride=2)
-      >>> # pool of non-square window
-      >>> m = MaxPool2d((3, 2), stride=(2, 1), channel_axis=-1)
-      >>> input = bst.random.randn(20, 50, 32, 16)
-      >>> output = m(input)
-      >>> output.shape
-      (20, 24, 31, 16)
-
-  .. _link:
-      https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  kernel_size: int, sequence of int
-    An integer, or a sequence of integers defining the window to reduce over.
-  stride: int, sequence of int
-    An integer, or a sequence of integers, representing the inter-window stride (default: `(1, ..., 1)`).
-  padding: str, int, sequence of tuple
-    Either the string `'SAME'`, the string `'VALID'`, or a sequence
-    of n `(low, high)` integer pairs that give the padding to apply before
-    and after each spatial dimension.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  mode: Mode
-    The computation mode.
-  name: optional, str
-    The object name.
-
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = None,
-      padding: Union[str, int, Tuple[int, ...], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(in_size=in_size,
-                     init_value=-jax.numpy.inf,
-                     computation=jax.lax.max,
-                     pool_dim=2,
-                     kernel_size=kernel_size,
-                     stride=stride,
-                     padding=padding,
-                     channel_axis=channel_axis,
-                     name=name, mode=mode)
-
-
-class MaxPool3d(_MaxPool):
-  r"""Applies a 3D max pooling over an input signal composed of several input planes.
-
-  In the simplest case, the output value of the layer with input size :math:`(N, D, H, W, C)`,
-  output :math:`(N, D_{out}, H_{out}, W_{out}, C)` and :attr:`kernel_size` :math:`(kD, kH, kW)`
-  can be precisely described as:
-
-  .. math::
-      \begin{aligned}
-        \text{out}(N_i, d, h, w) ={} & \max_{k=0, \ldots, kD-1} \max_{m=0, \ldots, kH-1} \max_{n=0, \ldots, kW-1} \\
-                                    & \text{input}(N_i, \text{stride[0]} \times d + k,
-                                    \text{stride[1]} \times h + m, \text{stride[2]} \times w + n, C_j)
-      \end{aligned}
-
-  If :attr:`padding` is non-zero, then the input is implicitly padded with negative infinity on both sides
-  for :attr:`padding` number of points. :attr:`dilation` controls the spacing between the kernel points.
-  It is harder to describe, but this `link`_ has a nice visualization of what :attr:`dilation` does.
-
-
-  Shape:
-      - Input: :math:`(N, D_{in}, H_{in}, W_{in}, C)` or :math:`(D_{in}, H_{in}, W_{in}, C)`.
-      - Output: :math:`(N, D_{out}, H_{out}, W_{out}, C)` or :math:`(D_{out}, H_{out}, W_{out}, C)`, where
-
-        .. math::
-            D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] - \text{dilation}[0] \times
-              (\text{kernel\_size}[0] - 1) - 1}{\text{stride}[0]} + 1\right\rfloor
-
-        .. math::
-            H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] - \text{dilation}[1] \times
-              (\text{kernel\_size}[1] - 1) - 1}{\text{stride}[1]} + 1\right\rfloor
-
-        .. math::
-            W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] - \text{dilation}[2] \times
-              (\text{kernel\_size}[2] - 1) - 1}{\text{stride}[2]} + 1\right\rfloor
-
-  Examples::
-
-      >>> import brainstate as bst
-      >>> # pool of square window of size=3, stride=2
-      >>> m = MaxPool3d(3, stride=2)
-      >>> # pool of non-square window
-      >>> m = MaxPool3d((3, 2, 2), stride=(2, 1, 2), channel_axis=-1)
-      >>> input = bst.random.randn(20, 50, 44, 31, 16)
-      >>> output = m(input)
-      >>> output.shape
-      (20, 24, 43, 15, 16)
-
-  .. _link:
-      https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  kernel_size: int, sequence of int
-    An integer, or a sequence of integers defining the window to reduce over.
-  stride: int, sequence of int
-    An integer, or a sequence of integers, representing the inter-window stride (default: `(1, ..., 1)`).
-  padding: str, int, sequence of tuple
-    Either the string `'SAME'`, the string `'VALID'`, or a sequence
-    of n `(low, high)` integer pairs that give the padding to apply before
-    and after each spatial dimension.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  mode: Mode
-    The computation mode.
-  name: optional, str
-    The object name.
-
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = None,
-      padding: Union[str, int, Tuple[int], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(in_size=in_size,
-                     init_value=-jax.numpy.inf,
-                     computation=jax.lax.max,
-                     pool_dim=3,
-                     kernel_size=kernel_size,
-                     stride=stride,
-                     padding=padding,
-                     channel_axis=channel_axis,
-                     name=name, mode=mode)
-
-
-class AvgPool1d(_AvgPool):
-  r"""Applies a 1D average pooling over an input signal composed of several input planes.
-
-  In the simplest case, the output value of the layer with input size :math:`(N, L, C)`,
-  output :math:`(N, L_{out}, C)` and :attr:`kernel_size` :math:`k`
-  can be precisely described as:
-
-  .. math::
-
-      \text{out}(N_i, l, C_j) = \frac{1}{k} \sum_{m=0}^{k-1}
-                             \text{input}(N_i, \text{stride} \times l + m, C_j)
-
-  If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides
-  for :attr:`padding` number of points.
-
-  Shape:
-      - Input: :math:`(N, L_{in}, C)` or :math:`(L_{in}, C)`.
-      - Output: :math:`(N, L_{out}, C)` or :math:`(L_{out}, C)`, where
-
-        .. math::
-            L_{out} = \left\lfloor \frac{L_{in} +
-            2 \times \text{padding} - \text{kernel\_size}}{\text{stride}} + 1\right\rfloor
-
-  Examples::
-
-      >>> import brainstate as bst
-      >>> # pool with window of size=3, stride=2
-      >>> m = AvgPool1d(3, stride=2)
-      >>> input = bst.random.randn(20, 50, 16)
-      >>> m(input).shape
-      (20, 24, 16)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  kernel_size: int, sequence of int
-    An integer, or a sequence of integers defining the window to reduce over.
-  stride: int, sequence of int
-    An integer, or a sequence of integers, representing the inter-window stride (default: `(1, ..., 1)`).
-  padding: str, int, sequence of tuple
-    Either the string `'SAME'`, the string `'VALID'`, or a sequence
-    of n `(low, high)` integer pairs that give the padding to apply before
-    and after each spatial dimension.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  mode: Mode
-    The computation mode.
-  name: optional, str
-    The object name.
-
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = 1,
-      padding: Union[str, int, Tuple[int, ...], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(in_size=in_size,
-                     init_value=0.,
-                     computation=jax.lax.add,
-                     pool_dim=1,
-                     kernel_size=kernel_size,
-                     stride=stride,
-                     padding=padding,
-                     channel_axis=channel_axis,
-                     name=name,
-                     mode=mode)
-
-
-class AvgPool2d(_AvgPool):
-  r"""Applies a 2D average pooling over an input signal composed of several input planes.
-
-  In the simplest case, the output value of the layer with input size :math:`(N, H, W, C)`,
-  output :math:`(N, H_{out}, W_{out}, C)` and :attr:`kernel_size` :math:`(kH, kW)`
-  can be precisely described as:
-
-  .. math::
-
-      out(N_i, h, w, C_j)  = \frac{1}{kH * kW} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1}
-                             input(N_i, stride[0] \times h + m, stride[1] \times w + n, C_j)
-
-  If :attr:`padding` is non-zero, then the input is implicitly zero-padded on both sides
-  for :attr:`padding` number of points.
-
-  Shape:
-      - Input: :math:`(N, H_{in}, W_{in}, C)` or :math:`(H_{in}, W_{in}, C)`.
-      - Output: :math:`(N, H_{out}, W_{out}, C)` or :math:`(H_{out}, W_{out}, C)`, where
-
-        .. math::
-            H_{out} = \left\lfloor\frac{H_{in}  + 2 \times \text{padding}[0] -
-              \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor
-
-        .. math::
-            W_{out} = \left\lfloor\frac{W_{in}  + 2 \times \text{padding}[1] -
-              \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor
-
-  Examples::
-
-      >>> import brainstate as bst
-      >>> # pool of square window of size=3, stride=2
-      >>> m = AvgPool2d(3, stride=2)
-      >>> # pool of non-square window
-      >>> m = AvgPool2d((3, 2), stride=(2, 1))
-      >>> input = bst.random.randn(20, 50, 32, , 16)
-      >>> output = m(input)
-      >>> output.shape
-      (20, 24, 31, 16)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  kernel_size: int, sequence of int
-    An integer, or a sequence of integers defining the window to reduce over.
-  stride: int, sequence of int
-    An integer, or a sequence of integers, representing the inter-window stride (default: `(1, ..., 1)`).
-  padding: str, int, sequence of tuple
-    Either the string `'SAME'`, the string `'VALID'`, or a sequence
-    of n `(low, high)` integer pairs that give the padding to apply before
-    and after each spatial dimension.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  mode: Mode
-    The computation mode.
-  name: optional, str
-    The object name.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = 1,
-      padding: Union[str, int, Tuple[int, ...], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(in_size=in_size,
-                     init_value=0.,
-                     computation=jax.lax.add,
-                     pool_dim=2,
-                     kernel_size=kernel_size,
-                     stride=stride,
-                     padding=padding,
-                     channel_axis=channel_axis,
-                     name=name,
-                     mode=mode)
-
-
-class AvgPool3d(_AvgPool):
-  r"""Applies a 3D average pooling over an input signal composed of several input planes.
-
-
-  In the simplest case, the output value of the layer with input size :math:`(N, D, H, W, C)`,
-  output :math:`(N, D_{out}, H_{out}, W_{out}, C)` and :attr:`kernel_size` :math:`(kD, kH, kW)`
-  can be precisely described as:
-
-  .. math::
-      \begin{aligned}
-          \text{out}(N_i, d, h, w, C_j) ={} & \sum_{k=0}^{kD-1} \sum_{m=0}^{kH-1} \sum_{n=0}^{kW-1} \\
-                                        & \frac{\text{input}(N_i, \text{stride}[0] \times d + k,
-                                                \text{stride}[1] \times h + m, \text{stride}[2] \times w + n, C_j)}
-                                               {kD \times kH \times kW}
-      \end{aligned}
-
-  If :attr:`padding` is non-zero, then the input is implicitly zero-padded on all three sides
-  for :attr:`padding` number of points.
-
-  Shape:
-      - Input: :math:`(N, D_{in}, H_{in}, W_{in}, C)` or :math:`(D_{in}, H_{in}, W_{in}, C)`.
-      - Output: :math:`(N, D_{out}, H_{out}, W_{out}, C)` or
-        :math:`(D_{out}, H_{out}, W_{out}, C)`, where
-
-        .. math::
-            D_{out} = \left\lfloor\frac{D_{in} + 2 \times \text{padding}[0] -
-                  \text{kernel\_size}[0]}{\text{stride}[0]} + 1\right\rfloor
-
-        .. math::
-            H_{out} = \left\lfloor\frac{H_{in} + 2 \times \text{padding}[1] -
-                  \text{kernel\_size}[1]}{\text{stride}[1]} + 1\right\rfloor
-
-        .. math::
-            W_{out} = \left\lfloor\frac{W_{in} + 2 \times \text{padding}[2] -
-                  \text{kernel\_size}[2]}{\text{stride}[2]} + 1\right\rfloor
-
-  Examples::
-
-      >>> import brainstate as bst
-      >>> # pool of square window of size=3, stride=2
-      >>> m = AvgPool3d(3, stride=2)
-      >>> # pool of non-square window
-      >>> m = AvgPool3d((3, 2, 2), stride=(2, 1, 2))
-      >>> input = bst.random.randn(20, 50, 44, 31, 16)
-      >>> output = m(input)
-      >>> output.shape
-      (20, 24, 43, 15, 16)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  kernel_size: int, sequence of int
-    An integer, or a sequence of integers defining the window to reduce over.
-  stride: int, sequence of int
-    An integer, or a sequence of integers, representing the inter-window stride (default: `(1, ..., 1)`).
-  padding: str, int, sequence of tuple
-    Either the string `'SAME'`, the string `'VALID'`, or a sequence
-    of n `(low, high)` integer pairs that give the padding to apply before
-    and after each spatial dimension.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  mode: Mode
-    The computation mode.
-  name: optional, str
-    The object name.
-
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      kernel_size: Size,
-      stride: Union[int, Sequence[int]] = 1,
-      padding: Union[str, int, Tuple[int, ...], Sequence[Tuple[int, int]]] = "VALID",
-      channel_axis: Optional[int] = -1,
-      mode: Mode = None,
-      name: Optional[str] = None,
-      in_size: Optional[Size] = None,
-  ):
-    super().__init__(in_size=in_size,
-                     init_value=0.,
-                     computation=jax.lax.add,
-                     pool_dim=3,
-                     kernel_size=kernel_size,
-                     stride=stride,
-                     padding=padding,
-                     channel_axis=channel_axis,
-                     name=name,
-                     mode=mode)
-
-
-def _adaptive_pool1d(x, target_size: int, operation: Callable):
-  """Adaptive pool 1D.
-
-  Args:
-    x: The input. Should be a JAX array of shape `(dim,)`.
-    target_size: The shape of the output after the pooling operation `(target_size,)`.
-    operation: The pooling operation to be performed on the input array.
-
-  Returns:
-    A JAX array of shape `(target_size, )`.
-  """
-  size = jnp.size(x)
-  num_head_arrays = size % target_size
-  num_block = size // target_size
-  if num_head_arrays != 0:
-    head_end_index = num_head_arrays * (num_block + 1)
-    heads = jax.vmap(operation)(x[:head_end_index].reshape(num_head_arrays, -1))
-    tails = jax.vmap(operation)(x[head_end_index:].reshape(-1, num_block))
-    outs = jnp.concatenate([heads, tails])
-  else:
-    outs = jax.vmap(operation)(x.reshape(-1, num_block))
-  return outs
-
-
-def _generate_vmap(fun: Callable, map_axes: List[int]):
-  map_axes = sorted(map_axes)
-  for axis in map_axes:
-    fun = jax.vmap(fun, in_axes=(axis, None, None), out_axes=axis)
-  return fun
-
-
-class _AdaptivePool(DnnLayer, ExplicitInOutSize):
-  """General N dimensional adaptive down-sampling to a target shape.
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  num_spatial_dims: int
-    The number of spatial dimensions.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  operation: Callable
-    The down-sampling operation.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-
-  def __init__(
-      self,
-      in_size: Size,
-      target_size: Size,
-      num_spatial_dims: int,
-      operation: Callable,
-      channel_axis: Optional[int] = -1,
-      name: Optional[str] = None,
-      mode: Optional[Mode] = None,
-  ):
-    super().__init__(name=name, mode=mode)
-
-    self.channel_axis = channel_axis
-    self.operation = operation
-    if isinstance(target_size, int):
-      self.target_shape = (target_size,) * num_spatial_dims
-    elif isinstance(target_size, Sequence) and (len(target_size) == num_spatial_dims):
-      self.target_shape = target_size
-    else:
-      raise ValueError("`target_size` must either be an int or tuple of length "
-                       f"{num_spatial_dims} containing ints.")
-
-    # in & out shapes
-    if in_size is not None:
-      in_size = tuple(in_size)
-      self.in_size = in_size
-      y = jax.eval_shape(self.update, jax.ShapeDtypeStruct((128,) + in_size, environ.dftype()))
-      self.out_size = y.shape[1:]
-
-  def update(self, x):
-    """Input-output mapping.
-
-    Parameters
-    ----------
-    x: Array
-      Inputs. Should be a JAX array of shape `(..., dim_1, dim_2, channels)`
-      or `(..., dim_1, dim_2)`.
-    """
-    # channel axis
-    channel_axis = self.channel_axis
-
-    if channel_axis:
-      if not 0 <= abs(channel_axis) < x.ndim:
-        raise ValueError(f"Invalid channel axis {channel_axis} for {x.shape}")
-      if channel_axis < 0:
-        channel_axis = x.ndim + channel_axis
-    # input dimension
-    if (x.ndim - (0 if channel_axis is None else 1)) < len(self.target_shape):
-      raise ValueError(f"Invalid input dimension. Except >={len(self.target_shape)} "
-                       f"dimensions (channel_axis={self.channel_axis}). "
-                       f"But got {x.ndim} dimensions.")
-    # pooling dimensions
-    pool_dims = list(range(x.ndim))
-    if channel_axis:
-      pool_dims.pop(channel_axis)
-
-    # pooling
-    for i, di in enumerate(pool_dims[-len(self.target_shape):]):
-      poo_axes = [j for j in range(x.ndim) if j != di]
-      op = _generate_vmap(_adaptive_pool1d, poo_axes)
-      x = op(x, self.target_shape[i], self.operation)
-    return x
-
-
-class AdaptiveAvgPool1d(_AdaptivePool):
-  r"""Applies a 1D adaptive max pooling over an input signal composed of several input planes.
-
-  The output size is :math:`L_{out}`, for any input size.
-  The number of output features is equal to the number of input planes.
-
-  Shape:
-      - Input: :math:`(N, L_{in}, C)` or :math:`(L_{in}, C)`.
-      - Output: :math:`(N, L_{out}, C)` or :math:`(L_{out}, C)`, where
-        :math:`L_{out}=\text{output\_size}`.
-
-  Examples:
-
-      >>> import brainstate as bst
-      >>> # target output size of 5
-      >>> m = AdaptiveMaxPool1d(5)
-      >>> input = bst.random.randn(1, 64, 8)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 5, 8)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self,
-               target_size: Union[int, Sequence[int]],
-               channel_axis: Optional[int] = -1,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = None,
-               in_size: Optional[Sequence[int]] = None, ):
-    super().__init__(in_size=in_size,
-                     target_size=target_size,
-                     channel_axis=channel_axis,
-                     num_spatial_dims=1,
-                     operation=jnp.mean,
-                     name=name,
-                     mode=mode)
-
-
-class AdaptiveAvgPool2d(_AdaptivePool):
-  r"""Applies a 2D adaptive max pooling over an input signal composed of several input planes.
-
-  The output is of size :math:`H_{out} \times W_{out}`, for any input size.
-  The number of output features is equal to the number of input planes.
-
-  Shape:
-      - Input: :math:`(N, H_{in}, W_{in}, C)` or :math:`(H_{in}, W_{in}, C)`.
-      - Output: :math:`(N, H_{out}, W_{out}, C)` or :math:`(H_{out}, W_{out}, C)`, where
-        :math:`(H_{out}, W_{out})=\text{output\_size}`.
-
-  Examples:
-
-      >>> import brainstate as bst
-      >>> # target output size of 5x7
-      >>> m = AdaptiveMaxPool2d((5, 7))
-      >>> input = bst.random.randn(1, 8, 9, 64)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 5, 7, 64)
-      >>> # target output size of 7x7 (square)
-      >>> m = AdaptiveMaxPool2d(7)
-      >>> input = bst.random.randn(1, 10, 9, 64)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 7, 7, 64)
-      >>> # target output size of 10x7
-      >>> m = AdaptiveMaxPool2d((None, 7))
-      >>> input = bst.random.randn(1, 10, 9, 64)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 10, 7, 64)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self,
-               target_size: Union[int, Sequence[int]],
-               channel_axis: Optional[int] = -1,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = None,
-               in_size: Optional[Sequence[int]] = None, ):
-    super().__init__(in_size=in_size,
-                     target_size=target_size,
-                     channel_axis=channel_axis,
-                     num_spatial_dims=2,
-                     operation=jnp.mean,
-                     name=name,
-                     mode=mode)
-
-
-class AdaptiveAvgPool3d(_AdaptivePool):
-  r"""Applies a 3D adaptive max pooling over an input signal composed of several input planes.
-
-  The output is of size :math:`D_{out} \times H_{out} \times W_{out}`, for any input size.
-  The number of output features is equal to the number of input planes.
-
-  Shape:
-      - Input: :math:`(N, D_{in}, H_{in}, W_{in}, C)` or :math:`(D_{in}, H_{in}, W_{in}, C)`.
-      - Output: :math:`(N, D_{out}, H_{out}, W_{out}, C)` or :math:`(D_{out}, H_{out}, W_{out}, C)`,
-        where :math:`(D_{out}, H_{out}, W_{out})=\text{output\_size}`.
-
-  Examples:
-
-      >>> import brainstate as bst
-      >>> # target output size of 5x7x9
-      >>> m = AdaptiveMaxPool3d((5, 7, 9))
-      >>> input = bst.random.randn(1, 8, 9, 10, 64)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 5, 7, 9, 64)
-      >>> # target output size of 7x7x7 (cube)
-      >>> m = AdaptiveMaxPool3d(7)
-      >>> input = bst.random.randn(1, 10, 9, 8, 64)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 7, 7, 7, 64)
-      >>> # target output size of 7x9x8
-      >>> m = AdaptiveMaxPool3d((7, None, None))
-      >>> input = bst.random.randn(1, 10, 9, 8, 64)
-      >>> output = m(input)
-      >>> output.shape
-      (1, 7, 9, 8, 64)
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self,
-               target_size: Union[int, Sequence[int]],
-               channel_axis: Optional[int] = -1,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = None,
-               in_size: Optional[Sequence[int]] = None, ):
-    super().__init__(in_size=in_size,
-                     target_size=target_size,
-                     channel_axis=channel_axis,
-                     num_spatial_dims=3,
-                     operation=jnp.mean,
-                     name=name,
-                     mode=mode)
-
-
-class AdaptiveMaxPool1d(_AdaptivePool):
-  """Adaptive one-dimensional maximum down-sampling.
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self,
-               target_size: Union[int, Sequence[int]],
-               channel_axis: Optional[int] = -1,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = None,
-               in_size: Optional[Sequence[int]] = None, ):
-    super().__init__(in_size=in_size,
-                     target_size=target_size,
-                     channel_axis=channel_axis,
-                     num_spatial_dims=1,
-                     operation=jnp.max,
-                     name=name,
-                     mode=mode)
-
-
-class AdaptiveMaxPool2d(_AdaptivePool):
-  """Adaptive two-dimensional maximum down-sampling.
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self,
-               target_size: Union[int, Sequence[int]],
-               channel_axis: Optional[int] = -1,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = None,
-               in_size: Optional[Sequence[int]] = None, ):
-    super().__init__(in_size=in_size,
-                     target_size=target_size,
-                     channel_axis=channel_axis,
-                     num_spatial_dims=2,
-                     operation=jnp.max,
-                     name=name,
-                     mode=mode)
-
-
-class AdaptiveMaxPool3d(_AdaptivePool):
-  """Adaptive three-dimensional maximum down-sampling.
-
-  Parameters
-  ----------
-  in_size: Sequence of int
-    The shape of the input tensor.
-  target_size: int, sequence of int
-    The target output shape.
-  channel_axis: int, optional
-    Axis of the spatial channels for which pooling is skipped.
-    If ``None``, there is no channel axis.
-  name: str
-    The class name.
-  mode: Mode
-    The computing mode.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self,
-               target_size: Union[int, Sequence[int]],
-               channel_axis: Optional[int] = -1,
-               name: Optional[str] = None,
-               mode: Optional[Mode] = None,
-               in_size: Optional[Sequence[int]] = None, ):
-    super().__init__(in_size=in_size,
-                     target_size=target_size,
-                     channel_axis=channel_axis,
-                     num_spatial_dims=3,
-                     operation=jnp.max,
-                     name=name,
-                     mode=mode)