MaxUnpool1d

Contents

MaxUnpool1d#

class brainstate.nn.MaxUnpool1d(kernel_size, stride=None, padding=0, channel_axis=-1, name=None, in_size=None)#

Computes a partial inverse of MaxPool1d.

MaxPool1d is not fully invertible, since the non-maximal values are lost. MaxUnpool1d takes in as input the output of MaxPool1d including the indices of the maximal values and computes a partial inverse in which all non-maximal values are set to zero.

Note

This function may produce nondeterministic gradients when given tensors on a CUDA device. See notes on reproducibility for more information.

Shape:

  • Input: \((N, L_{in}, C)\) or \((L_{in}, C)\)

  • Output: \((N, L_{out}, C)\) or \((L_{out}, C)\), where

    \[L_{out} = (L_{in} - 1) \times \text{stride} - 2 \times \text{padding} + \text{kernel\_size}\]

    or as given by output_size in the call operator

Parameters:

  • kernel_size (int | Sequence[int] | integer | Sequence[integer]) – Size of the max pooling window.

  • stride (int | Sequence[int]) – Stride of the max pooling window. Default: kernel_size

  • padding (int | Tuple[int, ...]) – Padding that was added to the input. Default: 0

  • channel_axis (int | None) – Axis of the channels. Default: -1

  • name (str | None) – Name of the module.

  • in_size (int | Sequence[int] | integer | Sequence[integer] | None) – Input size for shape inference.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>> # Create pooling and unpooling layers
>>> pool = MaxPool1d(2, stride=2, return_indices=True, channel_axis=-1)
>>> unpool = MaxUnpool1d(2, stride=2, channel_axis=-1)
>>> input = brainstate.random.randn(20, 50, 16)
>>> output, indices = pool(input)
>>> unpooled = unpool(output, indices)
>>> # unpooled will have shape (20, 100, 16) with zeros at non-maximal positions
Contents