N-D Pooling

Dimension-polymorphic pooling specs for spatial tensors and channels-first tensors.

Generic N-D pooling (channels-first, no batch)

These operators generalize the existing 2D pooling specs to an arbitrary spatial rank d.

Conventions:

• Input is channels-first: shape [C] ++ spatialDims.
• Pooling is applied independently per channel (like the existing 2D specs).
• kernel, stride, and padding are per-axis vectors (Vector Nat d).
• Padding is symmetric and uses zeros.

PyTorch comparisons (conceptual, without batch axis):

• max_pool_spec corresponds to torch.nn.functional.max_poolNd.
• avg_pool_spec corresponds to torch.nn.functional.avg_poolNd.

Layer configs + output shapes

structure Spec.MaxPoolSpec (d : ℕ) (kernel stride padding : Vector ℕ d) (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (hStride : ∀ (i : Fin d), stride.get i ≠ 0) : Type

Kernel/stride/padding configuration for N-D max pooling.

• kernelSizes : Vector ℕ d

  Kernel sizes per spatial axis (outermost to innermost).

• strideSizes : Vector ℕ d

  Strides per spatial axis (outermost to innermost).

• paddingSizes : Vector ℕ d

  Symmetric zero padding per spatial axis (outermost to innermost).

structure Spec.AvgPoolSpec (d : ℕ) (kernel stride padding : Vector ℕ d) (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (hStride : ∀ (i : Fin d), stride.get i ≠ 0) : Type

Kernel/stride/padding configuration for N-D average pooling.

• kernelSizes : Vector ℕ d

  Kernel sizes per spatial axis (outermost to innermost).

• strideSizes : Vector ℕ d

  Strides per spatial axis (outermost to innermost).

• paddingSizes : Vector ℕ d

  Symmetric zero padding per spatial axis (outermost to innermost).

def Spec.poolOutSpatial {d : ℕ} (inSpatial kernel stride : Vector ℕ d) : Vector ℕ d

"Valid" output spatial sizes (no padding): out = (in - k) / stride + 1 per axis.

def Spec.poolOutSpatialPad {d : ℕ} (inSpatial kernel stride padding : Vector ℕ d) : Vector ℕ d

Padded output spatial sizes: out = (in + 2*pad - k) / stride + 1 per axis.

def Spec.poolOutShape {d : ℕ} (inSpatial kernel stride : Vector ℕ d) : Shape

Output shape for single-channel N-D pooling (no padding).

def Spec.poolMultiOutShape {d : ℕ} (inC : ℕ) (inSpatial kernel stride : Vector ℕ d) : Shape

Output shape for channels-first N-D pooling (no padding; channels preserved).

def Spec.poolOutShapePad {d : ℕ} (inSpatial kernel stride padding : Vector ℕ d) : Shape

Output shape for single-channel N-D pooling with symmetric padding.

def Spec.poolMultiOutShapePad {d : ℕ} (inC : ℕ) (inSpatial kernel stride padding : Vector ℕ d) : Shape

Output shape for channels-first N-D pooling with symmetric padding (channels preserved).

def Spec.Private.tensorOfDims {α : Type} (dims : List ℕ) (f : List ℕ → α) : Tensor α (Shape.ofList dims)

def Spec.Private.foldlIndices' {β : Type} (dims : List ℕ) (init : β) (f : β → List ℕ → β) : β

def Spec.Private.paddedCoords? (outIdxs winIdxs stride : List ℕ) : Option (List ℕ)

def Spec.Private.unpadCoords? (padded padding : List ℕ) : Option (List ℕ)

def Spec.Private.coordsInBounds (idx dims : List ℕ) : Bool

def Spec.Private.getPaddedAverageInputVal {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (input : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs winIdxs stride padding : List ℕ) : α

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.

def Spec.Private.getPaddedMaxInputVal? {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (input : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs winIdxs stride padding : List ℕ) : Option α

Input lookup for hard max-pooling.

Unlike average-pooling, max-pooling should not insert a numeric zero for padded cells: PyTorch's max-pool semantics treat padding as -∞. TorchLean keeps the spec scalar-polymorphic by returning none for padded coordinates and letting the max fold ignore them.

def Spec.Private.kernelProd (kernel : List ℕ) : ℕ

def Spec.Private.maxPoolValue {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (input : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs kernel stride padding : List ℕ) : α

def Spec.Private.maxPoolJvpValue {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (input tangent : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs kernel stride padding : List ℕ) : α

Directional derivative of hard max-pooling for one N-D window.

The derivative is taken along the same winner selected by maxPoolValue. At ties we keep the first winner in row-major order, matching the VJP convention below and PyTorch's index convention.

def Spec.Private.avgPoolValue {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (input : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs kernel stride padding : List ℕ) : α

def Spec.Private.smoothMaxPoolValue {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (beta : α) (input : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs kernel stride padding : List ℕ) : α

def Spec.Private.smoothMaxPoolJvpValue {α : Type} [Context α] {d : ℕ} {inSpatial : Vector ℕ d} (beta : α) (input tangent : Tensor α (Shape.ofList inSpatial.toList)) (outIdxs kernel stride padding : List ℕ) : α

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 convention as smoothMaxPoolValue.

Forward (single-channel spatial tensor)

def Spec.maxPoolSpatialSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : MaxPoolSpec d kernel stride padding hKernel hStride) (input : Tensor α (Shape.ofList inSpatial.toList)) : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)

N-D max pooling on a spatial tensor (no explicit channel axis).

def Spec.maxPoolSpatialJvpSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : MaxPoolSpec d kernel stride padding hKernel hStride) (input tangent : Tensor α (Shape.ofList inSpatial.toList)) : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)

Forward-mode JVP for N-D hard max-pooling on a spatial tensor.

The derivative follows the same primal argmax as maxPoolSpatialSpec; at ties it keeps the first row-major maximizer. This is the correct directional derivative for TorchLean's chosen subgradient convention and matches the VJP tie policy.

def Spec.avgPoolSpatialSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : AvgPoolSpec d kernel stride padding hKernel hStride) (input : Tensor α (Shape.ofList inSpatial.toList)) : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)

N-D average pooling on a spatial tensor (no explicit channel axis).

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 (same as the 2D spec).
• 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=true behavior when padding is present).

def Spec.maxPoolSpatialBackwardSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (input : Tensor α (Shape.ofList inSpatial.toList)) (grad_output : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)) : Tensor α (Shape.ofList inSpatial.toList)

Backward/VJP for max_pool_spatial_spec.

Each output gradient is propagated to the argmax location in the corresponding input window. Ties keep the first position in row-major order.

def Spec.avgPoolSpatialBackwardSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : AvgPoolSpec d kernel stride padding hKernel hStride) (grad_output : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)) : Tensor α (Shape.ofList inSpatial.toList)

Backward/VJP for avg_pool_spatial_spec (single-channel).

Each output gradient is evenly distributed across its kernel window.

Forward (channels-first: C × spatial...)

def Spec.maxPoolSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (input : Tensor α (Shape.ofList (C :: inSpatial.toList))) : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))

N-D max pooling on a channels-first tensor: shape [C] ++ spatial.

def Spec.maxPoolJvpSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (input tangent : Tensor α (Shape.ofList (C :: inSpatial.toList))) : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))

N-D hard max-pool JVP on a channels-first tensor (channel-wise application).

def Spec.avgPoolSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : AvgPoolSpec d kernel stride padding hKernel hStride) (input : Tensor α (Shape.ofList (C :: inSpatial.toList))) : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))

N-D average pooling on a channels-first tensor: shape [C] ++ spatial.

Backward (channels-first: C × spatial...)

def Spec.maxPoolBackwardSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (input : Tensor α (Shape.ofList (C :: inSpatial.toList))) (grad_output : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))) : Tensor α (Shape.ofList (C :: inSpatial.toList))

Multi-channel VJP for max_pool_spec (apply spatial backward per channel).

def Spec.avgPoolBackwardSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : AvgPoolSpec d kernel stride padding hKernel hStride) (grad_output : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))) : Tensor α (Shape.ofList (C :: inSpatial.toList))

Multi-channel VJP for avg_pool_spec (apply spatial backward per channel).

Smooth max pooling (log-sum-exp surrogate)

def Spec.smoothMaxPoolSpatialSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : MaxPoolSpec d kernel stride padding hKernel hStride) (beta : α) (input : Tensor α (Shape.ofList inSpatial.toList)) : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)

Smooth log-sum-exp max pooling on a spatial tensor (no explicit channel axis).

def Spec.smoothMaxPoolSpatialJvpSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : MaxPoolSpec d kernel stride padding hKernel hStride) (beta : α) (input tangent : Tensor α (Shape.ofList inSpatial.toList)) : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)

Forward-mode JVP for N-D 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.

def Spec.smoothMaxPoolSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (beta : α) (input : Tensor α (Shape.ofList (C :: inSpatial.toList))) : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))

Smooth log-sum-exp max pooling on a channels-first tensor (channel-wise application).

def Spec.smoothMaxPoolJvpSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (beta : α) (input tangent : Tensor α (Shape.ofList (C :: inSpatial.toList))) : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))

N-D smooth max-pool JVP on a channels-first tensor (channel-wise application).

Smooth max pooling backward

def Spec.smoothMaxPoolSpatialBackwardSpec {α : Type} [Context α] {d : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (_layer : MaxPoolSpec d kernel stride padding hKernel hStride) (beta : α) (input : Tensor α (Shape.ofList inSpatial.toList)) (grad_output : Tensor α (Shape.ofList (poolOutSpatialPad inSpatial kernel stride padding).toList)) : Tensor α (Shape.ofList inSpatial.toList)

Backward/VJP for smooth_max_pool_spatial_spec (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ᵢ).

def Spec.smoothMaxPoolBackwardSpec {α : Type} [Context α] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} {hStride : ∀ (i : Fin d), stride.get i ≠ 0} (layer : MaxPoolSpec d kernel stride padding hKernel hStride) (beta : α) (input : Tensor α (Shape.ofList (C :: inSpatial.toList))) (grad_output : Tensor α (Shape.ofList (C :: (poolOutSpatialPad inSpatial kernel stride padding).toList))) : Tensor α (Shape.ofList (C :: inSpatial.toList))

Multi-channel VJP for smooth_max_pool_spec (apply spatial backward per channel).