GraphM Pooling Ops

N-dimensional and two-dimensional pooling builders with forward, JVP, and VJP payloads.

def Runtime.Autograd.Compiled.GraphM.maxPool {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : Var (Spec.Shape.ofList (C :: inSpatial.toList))) : MWith α Δ Γ (Var (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D max pooling (channels-first) on a single sample tensor (no batch axis).

PyTorch comparison: torch.nn.functional.max_pool1d / max_pool2d / max_pool3d depending on the spatial rank d.

Forward-mode status: implemented. The JVP follows the primal argmax selected by Spec.maxPoolJvpSpec, including the documented first-winner tie convention.

def Runtime.Autograd.Compiled.GraphM.avgPool {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (x : Var (Spec.Shape.ofList (C :: inSpatial.toList))) : MWith α Δ Γ (Var (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D average pooling (channels-first) on a single sample tensor (no batch axis).

PyTorch comparison: torch.nn.functional.avg_pool1d / avg_pool2d / avg_pool3d depending on the spatial rank d.

Forward-mode status: implemented. Average pooling is linear, so the JVP is the same average-pool map applied to the input tangent.

def Runtime.Autograd.Compiled.GraphM.smoothMaxPool {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (x : Var (Spec.Shape.ofList (C :: inSpatial.toList))) (beta : α) : MWith α Δ Γ (Var (Spec.Shape.ofList (C :: (Spec.poolOutSpatialPad inSpatial kernel stride padding).toList)))

N-D smooth max pooling (log-sum-exp surrogate) on a single sample tensor (no batch axis).

PyTorch comparison: there is no direct primitive; this is a differentiable approximation to max pooling.

Forward-mode status: implemented. The JVP is the softmax-weighted tangent of the log-sum-exp pooling window.

def Runtime.Autograd.Compiled.GraphM.maxPool2d {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Var (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

2D max-pooling (channel-first) on a single image tensor.

PyTorch comparison: torch.nn.functional.max_pool2d (without a batch dimension).

Forward-mode status: implemented. The JVP routes each output tangent through the argmax selected by the primal input.

def Runtime.Autograd.Compiled.GraphM.maxPool2dPad {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Var (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

2D max-pooling with explicit padding.

PyTorch comparison: torch.nn.functional.max_pool2d with padding.

Forward-mode status: implemented. Padding is fixed and the JVP follows the real primal winner, ignoring padded cells just like the forward pass.

def Runtime.Autograd.Compiled.GraphM.smoothMaxPool2d {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (x : Var (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) (beta : α) : MWith α Δ Γ (Var (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Smooth (soft) max-pooling, controlled by beta.

This is a differentiable approximation to max-pooling.

Forward-mode status: implemented. The JVP is the softmax-weighted tangent of the log-sum-exp pooling window.

def Runtime.Autograd.Compiled.GraphM.avgPool2d {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : Var (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim inC (Spec.Shape.dim ((inH - kH) / stride + 1) (Spec.Shape.dim ((inW - kW) / stride + 1) Spec.Shape.scalar))))

Average pooling (channel-first) on a single image tensor.

PyTorch comparison: torch.nn.functional.avg_pool2d (without a batch dimension).

Forward-mode status: implemented. Average pooling is linear, so the JVP is average pooling of the input tangent.

def Runtime.Autograd.Compiled.GraphM.avgPool2dPad {α Δ : Type} [Context α] [DecidableEq Spec.Shape] {Γ : List Spec.Shape} {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (x : Var (Spec.Shape.dim inC (Spec.Shape.dim inH (Spec.Shape.dim inW Spec.Shape.scalar)))) : MWith α Δ Γ (Var (Spec.Shape.dim inC (Spec.Shape.dim ((inH + 2 * padding - kH) / stride + 1) (Spec.Shape.dim ((inW + 2 * padding - kW) / stride + 1) Spec.Shape.scalar))))

Average pooling with explicit padding.

PyTorch comparison: torch.nn.functional.avg_pool2d with padding.

Forward-mode status: implemented. Padding is fixed and average pooling is linear, so the JVP is the padded average-pool map applied to the input tangent.