#23649·keras

torch pooling does two NHWC<->NCHW copies when one would do

Author: pctablet505Created Sep 16, 2026Updated Sep 16, 2026
Labelsbackend:torchlayers

On the torch backend, the pooling ops in keras/src/backend/torch/nn.py (max_pool, average_pool, adaptive_max_pool, adaptive_average_pool) call _transpose_spatial_inputs(inputs) in the default channels_last data format, which permutes NHWC->NCHW and then calls a plain .contiguous().

That .contiguous() is a full physical copy on every call. It is also unnecessary: a permuted NHWC->NCHW view of a contiguous input already has channels_last strides, so requesting torch.channels_last instead copies nothing. On top of that, the PyTorch ATen 2D pooling kernels (max_pool2d, avg_pool2d, adaptive_max_pool2d, adaptive_avg_pool2d) propagate channels_last to their output, so _transpose_spatial_outputs then returns an already-contiguous NHWC tensor rather than a non-contiguous permuted view.

This is exactly the same redundancy that #23279 identified for conv(), and it is not data-dependent: it fires on every pooling forward pass in channels_last.

_transpose_spatial_inputs already grew a channels_last_memory_format parameter in #23199 for the conv path; the pooling ops simply never opted in.

Measured on torch 2.11.0, input (32, 112, 112, 64), data_format="channels_last", outputs bit-exact (max diff = 0.0):

op CPU before CPU after MPS before MPS after
max_pool (valid) 9.37 ms 2.01 ms 2.21 ms 1.68 ms
average_pool (valid) 7.59 ms 1.18 ms 2.22 ms 1.11 ms
adaptive_max_pool 7.07 ms 2.03 ms 2.24 ms 1.63 ms
adaptive_average_pool 5.78 ms 1.18 ms 2.20 ms 1.11 ms

Fix: pass channels_last_memory_format=True to _transpose_spatial_inputs in the four pooling ops, matching what #23199 did for conv().

Part of the per-call Python-dispatch overhead series in #22561.