Reusable Neural-Network Blocks

Configuration records for MLP and convolution blocks, plus residual and branching combinators over already-built sequential models. The seeded block constructors live in NN.API.Seeded.

structure TorchLean.nn.MLP.Config : Type

MLP (multi-layer perceptron) configuration.

This builder produces a sequential stack of linear layers with activations and optional dropout.

PyTorch analogue: a hand-written nn.Sequential(Linear(...), ReLU(), ..., Linear(...)).

• hiddenWidths : List ℕ

  Hidden layer widths (each entry creates a Linear -> Activation stage).

• activation : Activation.Kind

  Activation used after each hidden linear layer.

• dropout? : Option Float

  Optional dropout probability after each activation.

def TorchLean.nn.MLP.Config.validate (config : Config) (inputWidth outputWidth : ℕ) : Except String Unit

Validate the input, output, hidden widths, and optional dropout probability before allocating parameter seeds.

structure TorchLean.nn.ConvBlock.Config (d : ℕ) : Type

Convolution followed by an activation and optional dropout.

• convolution : Convolution.Config d

  Convolution configuration.

• activation : Activation.Kind

  Activation applied after convolution.

• dropout? : Option Float

  Optional dropout probability applied after activation.

def TorchLean.nn.ConvBlock.Config.validate {d : ℕ} (config : Config d) (inputChannels : ℕ) (input : Tensor ℕ [d]) (kind : String := "ConvBlock") : Except String Unit

Validate convolution and dropout settings before allocating any parameter seeds.

structure TorchLean.nn.ConvPoolBlock.Config (d : ℕ) : Type

Convolution/activation followed by max pooling.

• block : ConvBlock.Config d

  Convolution and activation stage.

• pooling : Pooling.Config d

  Pooling stage.

def TorchLean.nn.ConvPoolBlock.Config.validate {d : ℕ} (config : Config d) (inputChannels : ℕ) (input : Tensor ℕ [d]) (kind : String := "ConvPoolBlock") : Except String Unit

Validate the convolution/dropout stage and the following pooling geometry.

def TorchLean.nn.activation {s : Shape} : Activation.Kind → Sequential s s

Interpret an activation kind as an elementwise sequential model.

def TorchLean.nn.Layer.residual {s : Shape} (inner : Sequential s s) : Layer s s

Residual (skip) connection as a single Layer.

Given inner : Seq s s this computes $x \mapsto \operatorname{inner}(x) + x$, the shape that appears in ResNet (He et al., Deep Residual Learning for Image Recognition, CVPR 2016) and in every Transformer sublayer since Vaswani et al., Attention Is All You Need (NeurIPS 2017).

The Layer namespace holds the raw single-layer constructors; the same name without the namespace returns a Sequential. That split is why there is no residualLayer: the return type belongs in the namespace, not in the identifier.

def TorchLean.nn.residual {s : Shape} (inner : Sequential s s) : Sequential s s

The same residual connection, wrapped as a one-element Sequential model.

Example:

def inner : nn.Builder (nn.Sequential [64] [64]) :=
  nn.Sequential![nn.linear 64 64, nn.relu]

-- `x + inner x`, the connection that made deep stacks trainable.
def model : nn.Builder (nn.Sequential [64] [64]) := do
  pure (nn.residual (← inner))

Branching (skip connections)

Seq is linear, but skip connections require two computations to consume the same input. The generic constructor below owns the shared state plumbing; public blocks choose how to combine the two outputs.

def TorchLean.nn.Layer.combineBranches {σ τ₁ τ₂ υ : Shape} (kind : String) (f : Sequential σ τ₁) (g : Sequential σ τ₂) (combine : {α : Type} → [inst : Storage α] → [inst_1 : Context α] → {m : Type → Type} → [inst_2 : Monad m] → [inst_3 : Runtime.Autograd.Torch.Ops m α] → Runtime.ValueRef m α τ₁ → Runtime.ValueRef m α τ₂ → m (Runtime.ValueRef m α υ)) : Layer σ υ

Run two sequential branches on the same input and combine their outputs.

Parameters and persistent buffers are stored as state(f) ++ state(g). The combining operation is polymorphic in the runtime, so eager execution and graph lowering share the same branch structure.

def TorchLean.nn.Layer.addBranches {σ τ : Shape} (f g : Sequential σ τ) : Layer σ τ

Two branches over one input, outputs added. The Sequential form is nn.addBranches.

def TorchLean.nn.addBranches {σ τ : Shape} (f g : Sequential σ τ) : Sequential σ τ

Combine two models with the same input/output shapes by summing their outputs.

This is a typed residual-add block: addBranches f g represents the model $x \mapsto f(x) + g(x)$, and its parameter list is the concatenation of the two branches’ parameter lists.

Example:

-- Two branches over the same input, outputs summed. Parameters are stored as the first branch's
-- list followed by the second's.
def model : nn.Builder (nn.Sequential [32] [8]) := do
  let wide ← nn.Sequential![nn.linear 32 8, nn.relu]
  let shortcut ← nn.linear 32 8
  pure (nn.addBranches wide shortcut)

def TorchLean.nn.Layer.concatBranches {σ s : Shape} {n m : ℕ} (f : Sequential σ (s.prependDim n)) (g : Sequential σ (s.prependDim m)) : Layer σ (s.prependDim (n + m))

Concatenate two branch outputs along their first axis.

Both branches consume the same input. Their outputs must agree on every remaining dimension, and the result records the sum of their first-axis extents in its type. This is the general typed skip connection needed by encoder-decoder models; arbitrary outer batch axes can be added with mapLeading.

def TorchLean.nn.concatBranches {σ s : Shape} {n m : ℕ} (f : Sequential σ (s.prependDim n)) (g : Sequential σ (s.prependDim m)) : Sequential σ (s.prependDim (n + m))

Concatenate two typed branches along their first output axis.

Example:

-- Concatenation along the first output axis: `[4, 8]` next to `[6, 8]` gives `[10, 8]`, and the
-- addition happens in the type rather than in a runtime shape check.
def model : nn.Builder (nn.Sequential [16] [10, 8]) := do
  let left ← nn.Sequential![nn.linear 16 (4 * 8), nn.reshape [4 * 8] [4, 8]]
  let right ← nn.Sequential![nn.linear 16 (6 * 8), nn.reshape [6 * 8] [6, 8]]
  pure (nn.concatBranches left right)