Spatial Pooling #
Dimension-polymorphic pooling specs for spatial tensors and channels-first tensors.
Arbitrary-rank pooling (channels-first, no batch) #
These operators define pooling over an arbitrary spatial rank d.
Conventions:
- Input is channels-first: shape
[C] ++ spatialDims. - Pooling is applied independently per channel.
kernel,stride, andpaddingare per-axis vectors (Tensor Nat [d]).- Padding is symmetric. Average pooling counts padded positions as zeros. Max pooling ignores padded positions; a window with no input position is outside PyTorch's valid max-pool domain and is totalized to zero by the scalar-polymorphic TorchLean spec.
PyTorch comparisons (conceptual, without batch axis):
maxPoolSpeccorresponds to PyTorch's rank-specificmax_pool1d,max_pool2d, andmax_pool3doperations.avgPoolSpeccorresponds to PyTorch's rank-specificavg_pool1d,avg_pool2d, andavg_pool3doperations.
Layer configs + output shapes #
Witness that an arbitrary-rank max-pooling configuration has nonzero kernel and stride on every axis.
The structure has no fields: it is a phantom type whose indices force every construction site to
supply the nonzero-kernel and nonzero-stride proofs, and it is the value threaded through the
runtime, lowering, and verification code paths. The pooling specs themselves do not eliminate
these proofs. A zero kernel or stride is instead excluded at the shape level: poolOutDim returns
an empty output axis for a zero kernel, so no window is ever evaluated, and the getD 0 fallback
inside Pooling.Internal.maxPoolValue is unreachable for the shapes this witness admits.
The tensors are indices of the type rather than duplicate structure fields, so a value cannot advertise one configuration while its type describes another.
Instances For
Witness that an arbitrary-rank average-pooling configuration has nonzero kernel and stride.
Like MaxPoolSpec, it is a field-free phantom type; see that docstring for what the proofs do.
Instances For
Output spatial sizes with symmetric padding.
Pooling follows the usual floor-mode sliding-window formula, but an empty input axis, empty kernel, or padding larger than half the kernel gives an empty output axis. The last condition is part of the pooling contract used by PyTorch and by TorchLean's native implementations; it is not a restriction on convolution.
Instances For
Output spatial sizes without padding.
An invalid axis (empty input, zero kernel, zero stride, or a kernel larger than the input) has size zero.
Instances For
Apply poolOutDim independently to each spatial axis.
Instances For
Pooling over the complete spatial extent produces one value on every spatial axis.
Output shape for single-channel arbitrary-rank pooling (no padding).
Instances For
Output shape for channels-first arbitrary-rank pooling (no padding; channels preserved).
Instances For
Output shape for single-channel arbitrary-rank pooling with symmetric padding.
Instances For
Output shape for channels-first arbitrary-rank pooling with symmetric padding (channels preserved).
Instances For
Choose the input-space pivot whose scaled value is maximal.
Instances For
Fold over every coordinate of a rectangular index box given by dims.
Pooling walks a window whose rank is only known at runtime, so the walk is a fold over a list of
extents rather than nested Fin loops. The accumulator sees each coordinate list once, in row-major
order.
Instances For
Turn a padded-input coordinate back into a real input coordinate.
Returns none when the coordinate lands inside the padding, which is how the callers tell a genuine
input cell from a fabricated one; average and max pooling then treat that answer differently.
Instances For
Whether a runtime coordinate list is in range for dims, and of the right length.
Instances For
Input lookup for average/smooth pooling.
For average-style pooling, padded cells contribute numeric zero and are still counted by the
denominator chosen by the surrounding pooling spec. We keep this separate from
getPaddedMaxInputVal?, where padded cells must be ignored rather than treated as zero.
Instances For
Input lookup for hard max-pooling.
Unlike average pooling, max pooling does not insert numeric zero for an individual padded cell:
PyTorch's valid max-pool configurations behave as though those cells were -∞. TorchLean keeps
the spec scalar-polymorphic by returning none for padded coordinates and ignoring them in the
max fold. poolOutSpatialPad rejects empty input axes, empty kernels, and padding beyond PyTorch's
half-kernel restriction, so every emitted output window contains at least one input coordinate.
Instances For
Number of cells in a pooling window, the denominator average pooling counts with.
Instances For
Start of adaptive-pooling bin i: floor(i * input / output).
Instances For
End of adaptive-pooling bin i: ceil((i + 1) * input / output).
Instances For
Every adaptive-pooling bin is nonempty when the input and output extents are positive.
Both hypotheses are needed: an empty input axis gives start = end = 0, and output = 0 makes
both quotients zero under Lean's total division. This is the fact that justifies the hInput and
hOutput arguments of adaptiveAvgPoolSpec and adaptiveMaxPoolSpec: the average never divides
by zero and the max fold always sees at least one element.
Per-axis bin extents for one adaptive-pooling output coordinate.
Adaptive pooling is the case where the window is not fixed: bin widths vary with the coordinate when the input extent is not a multiple of the output extent, which is exactly why the walk is over a computed extent list.
Instances For
Average of one adaptive-pooling bin.
adaptiveStart_lt_adaptiveEnd is what keeps this honest: with positive input and output extents the
bin is nonempty, so the division below is never a division by zero.
Instances For
Maximum over one adaptive-pooling bin.
The getD 0 fallback would only be reached on an empty bin, which adaptiveStart_lt_adaptiveEnd
rules out for the shapes the public specs accept.
Instances For
Maximum over one fixed pooling window, ignoring padded cells.
Padded cells are skipped rather than treated as zero, since a zero would win the max over an
all-negative window. poolOutSpatialPad guarantees at least one real cell per window.
Instances For
Selected-branch tangent for one hard max-pooling window.
The tangent follows the same winner selected by maxPoolValue. At a tie this is a deterministic
generalized-derivative convention, not the mathematical directional derivative of max.
Instances For
Average over one fixed pooling window, counting padded cells as zero.
This is PyTorch's count_include_pad = True behaviour: the denominator is the whole window, so a
padded edge pulls the average toward zero rather than being renormalized away.
Instances For
Input-space pivot whose scaled value is maximal over a nonempty arbitrary-rank pooling window.
Instances For
Evaluate one arbitrary-rank smooth-max window with the sign-aware input-space pivot.
Instances For
Directional derivative of the smooth log-sum-exp pooling value.
For y = beta⁻¹ log Σ exp(beta*xᵢ), the directional derivative is
Σ softmax(beta*xᵢ) * dxᵢ, using the same zero-padding and stable input-pivot convention as
smoothMaxPoolValue.
Instances For
Forward (single-channel spatial tensor) #
arbitrary-rank max pooling on a spatial tensor (no explicit channel axis).
Instances For
Selected-branch linearization for arbitrary-rank hard max-pooling on a spatial tensor.
Away from ties this is the ordinary JVP. At ties it follows the first row-major primal maximizer, matching the VJP convention but not claiming an analytic directional derivative.
Instances For
arbitrary-rank average pooling on a spatial tensor (no explicit channel axis).
Instances For
Backward (single-channel spatial tensor) #
These are the VJPs of the forward pooling specs above.
Conventions:
- For max pooling, ties are broken by first occurrence in row-major order.
- For max pooling, padded cells are ignored, modeling PyTorch's
-∞padding without requiring a scalar-polymorphic infinity constant. - For average pooling, gradients are evenly distributed across the full kernel window
(
count_include_pad=truebehavior when padding is present).
Backward/VJP for maxPoolSpatialSpec.
Each output gradient is propagated to the argmax location in the corresponding input window. Ties keep the first position in row-major order.
Instances For
Backward/VJP for avgPoolSpatialSpec (single-channel).
Each output gradient is evenly distributed across its kernel window.
Instances For
Forward (channels-first: C × spatial...) #
arbitrary-rank max pooling on a channels-first tensor: shape [C] ++ spatial.
Instances For
arbitrary-rank hard max-pool selected-branch linearization, applied channel-wise.
Instances For
arbitrary-rank average pooling on a channels-first tensor: shape [C] ++ spatial.
Instances For
Adaptive pooling #
Adaptive pooling partitions every input axis into a requested number of bins. Unlike fixed-window pooling, its window sizes depend on the output index. The same definition handles sequence, image, volume, and higher-rank tensors.
Adaptive average pooling on a channels-first tensor of arbitrary spatial rank.
The two hypotheses make every bin nonempty (Pooling.Internal.adaptiveStart_lt_adaptiveEnd), so
the per-bin average never divides by zero. They are not eliminated inside the definition.
Instances For
Adaptive max pooling on a channels-first tensor of arbitrary spatial rank.
The two hypotheses make every bin nonempty (Pooling.Internal.adaptiveStart_lt_adaptiveEnd), so
the getD 0 fallback in Pooling.Internal.adaptiveMaxPoolValue is never taken.
Instances For
Backward (channels-first: C × spatial...) #
Multi-channel VJP for maxPoolSpec (apply spatial backward per channel).
Instances For
Multi-channel VJP for avgPoolSpec (apply spatial backward per channel).
Instances For
Smooth max pooling (log-sum-exp surrogate) #
Smooth log-sum-exp max pooling on a spatial tensor (no explicit channel axis).
The temperature parameter must be nonzero because the forward expression contains 1 / beta.
Instances For
Forward-mode JVP for arbitrary-rank smooth max-pooling on a spatial tensor.
For the log-sum-exp surrogate this is the softmax-weighted sum of the input tangent over each
window. It is the forward-mode counterpart of smoothMaxPoolSpatialBackwardSpec.
Instances For
Smooth log-sum-exp max pooling on a channels-first tensor (channel-wise application).
The temperature parameter must be nonzero because the forward expression contains 1 / beta.
Instances For
arbitrary-rank smooth max-pool JVP on a channels-first tensor (channel-wise application).
Instances For
Smooth max pooling backward #
Backward/VJP for smoothMaxPoolSpatialSpec (log-sum-exp surrogate).
For a window x₁,…,xₙ, the surrogate is:
y = (1/beta) * log(∑ exp(beta*xᵢ))
and the VJP distributes upstream gradient proportionally to exp(beta*xᵢ).
The implementation evaluates the equivalent max/min-shifted weights so large finite inputs do not
overflow before normalization.
Instances For
Multi-channel VJP for smoothMaxPoolSpec (apply spatial backward per channel).