Core Tape Convolution and Pooling

This file implements the pure tape nodes for convolution, transposed convolution, and pooling. These nodes are backend-independent: they record forward values, parents, and backward closures using the spec-layer definitions before CUDA or compiled backends enter the picture.

def Runtime.Autograd.Tape.conv {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} (t : Tape α) (kernelId biasId inputId : ℕ) (name : String := "conv") : Result (Tape α × ℕ)

N-D convolution for channels-first tensors (inC, spatial...) (no batch axis).

This is the generic counterpart to conv2d; conv2d is implemented as a specialization with d = 2, scalar stride, and scalar padding.

def Runtime.Autograd.Tape.conv2d {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (t : Tape α) (kernelId biasId inputId : ℕ) : Result (Tape α × ℕ)

2D convolution for channel-first images (inC,inH,inW) (no batch axis).

PyTorch comparison: torch.nn.functional.conv2d specialized to a single image.

def Runtime.Autograd.Tape.convTranspose {α : Type} [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] [DecidableEq Spec.Shape] {d inC outC : ℕ} {kernel stride padding inSpatial : Vector ℕ d} (t : Tape α) (kernelId biasId inputId : ℕ) (name : String := "conv_transpose") : Result (Tape α × ℕ)

N-D transpose convolution for channels-first tensors (inC, spatial...) (no batch axis).

This is the generic counterpart to conv_transpose2d.

Kernel layout matches the spec/PyTorch convention (inC, outC, kernel[0], ..., kernel[d-1]).

PyTorch comparison: torch.nn.functional.conv_transpose{d}d specialized to a single sample (no batch axis).

def Runtime.Autograd.Tape.convTranspose2d {α : Type} [Context α] [DecidableEq Spec.Shape] {inC outC kH kW stride padding inH inW : ℕ} {h1 : inC ≠ 0} {h2 : kH ≠ 0} {h3 : kW ≠ 0} (t : Tape α) (kernelId biasId inputId : ℕ) : Result (Tape α × ℕ)

2D transpose convolution for channel-first images (inC,inH,inW) (no batch axis).

This is implemented as a specialization of conv_transpose with d = 2, scalar stride, and scalar padding. Kernel layout matches the spec/PyTorch convention (inC,outC,kH,kW).

PyTorch comparison: torch.nn.functional.conv_transpose2d specialized to a single image.

def Runtime.Autograd.Tape.maxPool {α : Type} [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

N-D max pooling for channels-first tensors (C, spatial...) (no batch axis).

Padding is symmetric per-axis and uses zeros. To model unpadded pooling, pass padding := 0 on every axis.

def Runtime.Autograd.Tape.avgPool {α : Type} [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} (hKernel : ∀ (i : Fin d), kernel.get i ≠ 0) (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

N-D average pooling for channels-first tensors (C, spatial...) (no batch axis).

Padding is symmetric per-axis and uses zeros; pooling uses count_include_pad=true semantics.

def Runtime.Autograd.Tape.smoothMaxPool {α : Type} [Context α] [DecidableEq Spec.Shape] {d C : ℕ} {inSpatial kernel stride padding : Vector ℕ d} {hKernel : ∀ (i : Fin d), kernel.get i ≠ 0} (t : Tape α) (xId : ℕ) (beta : α) : Result (Tape α × ℕ)

N-D smooth max pooling (log-sum-exp surrogate) for channels-first tensors (C, spatial...).

def Runtime.Autograd.Tape.maxPool2d {α : Type} [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

2D max-pooling for channel-first images (no batch axis).

PyTorch comparison: torch.nn.functional.max_pool2d.

def Runtime.Autograd.Tape.maxPool2dPad {α : Type} [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

2D max-pooling with padding for channel-first images (no batch axis).

PyTorch comparison: max_pool2d(..., padding=...).

def Runtime.Autograd.Tape.smoothMaxPool2d {α : Type} [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} {h1 : kH ≠ 0} {h2 : kW ≠ 0} (t : Tape α) (xId : ℕ) (beta : α) : Result (Tape α × ℕ)

Smooth approximation of max-pooling (softmax pooling).

This is not a standard PyTorch primitive; it is useful for differentiable relaxations.

def Runtime.Autograd.Tape.avgPool2d {α : Type} [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

2D average-pooling for channel-first images (no batch axis).

PyTorch comparison: torch.nn.functional.avg_pool2d.

def Runtime.Autograd.Tape.avgPool2dPad {α : Type} [Context α] [DecidableEq Spec.Shape] {kH kW inH inW inC stride padding : ℕ} (h1 : kH ≠ 0) (h2 : kW ≠ 0) (t : Tape α) (xId : ℕ) : Result (Tape α × ℕ)

2D average-pooling with padding for channel-first images (no batch axis).

PyTorch comparison: avg_pool2d(..., padding=...).