Linear module wrappers

NN/Spec/Layers/Linear.lean defines the math-level linear layer spec (parameters + gradients).

This file provides NNModuleSpec wrappers so linear layers can be:

• composed in a SpecChain with compile-time shape checking, and
• annotated with lightweight export metadata (PyTorch pretty-printing).

These wrappers are thin on purpose: the meaning is still the underlying linear_spec; the extra fields are just metadata.

Most readers can map these directly to PyTorch: LinearModuleSpec is nn.Linear, and the sequence helpers are the common pattern "apply a linear layer to each timestep, then maybe select the last step for classification".

def Spec.LinearModuleSpec {α : Type} [Add α] [Mul α] [Zero α] {inDim outDim : ℕ} (m : LinearSpec α inDim outDim) : ModSpec.NNModuleSpec α (Shape.dim inDim Shape.scalar) (Shape.dim outDim Shape.scalar)

Wrap LinearSpec forward as an NNModuleSpec.

def Spec.LinearSeqModuleSpec {α : Type} [Context α] {seqLen hiddenSize outputSize : ℕ} (m : LinearSpec α hiddenSize outputSize) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim hiddenSize Shape.scalar)) (Shape.dim seqLen (Shape.dim outputSize Shape.scalar))

Apply a linear layer independently at each timestep (sequence lift).

def Spec.getLastTimestepSpec {α : Type} [Context α] {seqLen hiddenSize : ℕ} (seq : Tensor α (Shape.dim seqLen (Shape.dim hiddenSize Shape.scalar))) (h : seqLen ≠ 0) : Tensor α (Shape.dim hiddenSize Shape.scalar)

Extract the last timestep of a sequence (requires seqLen ≠ 0).

In PyTorch terms: seq[-1] when seq has shape (seqLen, hiddenSize).

def Spec.LinearClassifierModuleSpec {α : Type} [Context α] {seqLen hiddenSize numClasses : ℕ} (m : LinearSpec α hiddenSize numClasses) (h : seqLen ≠ 0) : ModSpec.NNModuleSpec α (Shape.dim seqLen (Shape.dim hiddenSize Shape.scalar)) (Shape.dim numClasses Shape.scalar)

Sequence classifier: take last timestep, then apply a linear projection to numClasses.

def Spec.twoLayerLinear {α : Type} [Add α] [Mul α] [Zero α] {inDim hidDim outDim : ℕ} (l1 : LinearSpec α inDim hidDim) (l2 : LinearSpec α hidDim outDim) : ModSpec.SpecChain α (Shape.dim inDim Shape.scalar) (Shape.dim outDim Shape.scalar)

Example: Linear → Linear composition via SpecChain.