Padded 2D Pooling

Symmetric-padding variants of the 2D pooling specs and their backward maps.

Padded pooling (symmetric padding)

For max-pooling, padded locations are not real input elements and are ignored when selecting the maximum. This is the scalar-polymorphic way to model PyTorch's -∞ max-pool padding without adding a backend-specific infinity constant to Context α.

For average pooling, this corresponds to including padded zeros in the average (PyTorch's default count_include_pad = true).

def Spec.unpadImage {α : Type} [Context α] {H W padding : ℕ} (img : Image (H + 2 * padding) (W + 2 * padding) α) : Image H W α

Remove symmetric zero-padding from a single-channel image.

def Spec.unpadMultiChannel {α : Type} [Context α] {C H W padding : ℕ} (img : MultiChannelImage C (H + 2 * padding) (W + 2 * padding) α) : MultiChannelImage C H W α

Remove symmetric zero-padding from a multi-channel image (channel-wise unpad_image).

def Spec.maxPool2dMultiSpecPad {α : Type} [Context α] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} {hStride : stride ≠ 0} (layer : MaxPool2DSpec kH kW stride h1 h2 hStride) (input : MultiChannelImage inC inH inW α) : MultiChannelImage inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1) α

Multi-channel max-pooling forward pass with PyTorch-style padding (-∞ outside bounds).

def Spec.maxPool2dMultiJvpSpecPad {α : Type} [Context α] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} {hStride : stride ≠ 0} (layer : MaxPool2DSpec kH kW stride h1 h2 hStride) (input tangent : MultiChannelImage inC inH inW α) : MultiChannelImage inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1) α

Forward-mode JVP for padded hard max-pooling.

Padding cells are ignored exactly as in maxPool2dMultiSpecPad, so the tangent is taken from the primal winner among real input locations only. If a window contains no real input cells, the forward value and tangent are both 0.

def Spec.avgPool2dMultiSpecPad {α : Type} [Context α] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) {hStride : stride ≠ 0} (layer : AvgPool2DSpec kH kW stride h1 h2 hStride) (input : MultiChannelImage inC inH inW α) : MultiChannelImage inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1) α

Multi-channel average pooling forward pass with symmetric zero padding.

def Spec.maxPool2dMultiBackwardSpecPad {α : Type} [Context α] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} {hStride : stride ≠ 0} (layer : MaxPool2DSpec kH kW stride h1 h2 hStride) (input : MultiChannelImage inC inH inW α) (grad_output : MultiChannelImage inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1) α) : MultiChannelImage inC inH inW α

Multi-channel max-pooling backward pass with PyTorch-style padding (-∞ outside bounds).

def Spec.avgPool2dMultiBackwardSpecPad {α : Type} [Context α] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) {hStride : stride ≠ 0} (layer : AvgPool2DSpec kH kW stride h1 h2 hStride) (grad_output : MultiChannelImage inC ((inH + 2 * padding - kH) / stride + 1) ((inW + 2 * padding - kW) / stride + 1) α) : MultiChannelImage inC inH inW α

Multi-channel average-pooling backward pass with symmetric padding (backprop then unpad).

def Spec.smoothMaxPool2dBackwardSpec {α : Type} [Context α] {kH kW inH inW stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} {hStride : stride ≠ 0} (_layer : MaxPool2DSpec kH kW stride h1 h2 hStride) (beta : α) (input : Image inH inW α) (grad_output : Image ((inH - kH) / stride + 1) ((inW - kW) / stride + 1) α) : Image inH inW α

Backward/VJP for smooth_max_pool2d_spec (log-sum-exp surrogate).

def Spec.smoothMaxPool2dMultiBackwardSpec {α : Type} [Context α] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} {hStride : stride ≠ 0} (layer : MaxPool2DSpec kH kW stride h1 h2 hStride) (beta : α) (input : MultiChannelImage inC inH inW α) (grad_output : MultiChannelImage inC ((inH - kH) / stride + 1) ((inW - kW) / stride + 1) α) : MultiChannelImage inC inH inW α

Multi-channel backward for smooth_max_pool2d_multi_spec (apply per channel).