mlx.nn.MaxPool1d#

class mlx.nn.MaxPool1d(kernel_size: Union[int, Tuple[int, int]], stride: Optional[Union[int, Tuple[int, int]]] = None, padding: Optional[Union[int, Tuple[int, int]]] = 0)#

Applies 1-dimensional max pooling.

Assuming an input of shape \((N, L, C)\) and kernel_size is \(k\), the output is a tensor of shape \((N, L_{out}, C)\), given by:

\[\text{out}(N_i, t, C_j) = \max_{m=0, \ldots, k - 1} \text{input}(N_i, \text{stride} \times t + m, C_j),\]

where \(L_{out} = \left\lfloor \frac{L + 2 \times \text{padding} - \text{kernel_size}}{\text{stride}}\right\rfloor + 1\).

Parameters:

  • kernel_size (int or tuple(int)) – The size of the pooling window kernel.

  • stride (int or tuple(int), optional) – The stride of the pooling window. Default: kernel_size.

  • padding (int or tuple(int), optional) – How much negative infinity padding to apply to the input. The padding amount is applied to both sides of the spatial axis. Default: 0.

Examples

>>> import mlx.core as mx
>>> import mlx.nn.layers as nn
>>> x = mx.random.normal(shape=(4, 16, 5))
>>> pool = nn.MaxPool1d(kernel_size=2, stride=2)
>>> pool(x)