Models over Leading Dimensions

This module lifts sequential models over any number of leading tensor dimensions. A batch is the common case leading = [batch]; shapes such as [batch, time] use the same machinery.

mapLeading applies a model separately at every leading index. The implementation module behind the layer constructors also flattens leading dimensions for layers that already accept one outer dimension; keeping that distinction explicit matters for stateful layers, whose buffer updates may depend on whether the leading positions are processed together or one at a time.

partial def TorchLean.nn.mapLeading (leading : Shape) {σ τ : Shape} : Sequential σ τ → Sequential (leading.concat σ) (leading.concat τ)

Apply a sequential model separately at every index of leading.

All positions use the same model parameters. Buffer updates are evaluated in lexicographic order over the leading indices.

Example:

def perSample : nn.Builder (nn.Sequential [2] [1]) :=
  nn.Sequential![nn.linear 2 8, nn.relu, nn.linear 8 1]

-- One model, applied at each of five positions of a new leading axis, sharing its parameters.
def model : nn.Builder (nn.Sequential [5, 2] [5, 1]) := do
  pure (nn.mapLeading [5] (← perSample))