MaxPool

Pools the input by taking the maximum of a window slice.

Abstract Signature:

MaxPool(inputs: Tensor, window_shape: Union[int, Tuple[int, ...]], strides: Union[int, Tuple[int, ...]], padding: str = VALID)

PyTorchJAX (Core)Flax NNX

PyTorch

API: torch.nn.functional.max_pool2d

Strategy: Direct Mapping

Official Docs ↗

JAX (Core)

API: jax.lax.reduce_window

Strategy: Plugin (pooling_shim)

Flax NNX

API: flax.nnx.max_pool

Strategy: Direct Mapping

Official Docs ↗