Diagonal S4-style state-space layer

This module provides TorchLean's diagonal recurrent SSM layer in the S4 family. It exposes the state-space recurrence used by S4-style models:

h_{t+1} = A h_t + B x_t, y_t = C h_{t+1} + D x_t.

The diagonal form is intentional: it shares the selective-scan core used by Mamba-style models, admits direct recurrence proofs, and can be connected to convolutional S4 kernels through a separate structured-kernel layer.

Reference: Gu, Goel, Ré. "Efficiently Modeling Long Sequences with Structured State Spaces", ICLR 2022.

Implementation status

No API builder implements this layer; the nn.* builders have no S4 layer. The relating theorem is diagonalS4_runArray_append_outputs_prefix in NN/MLTheory/Proofs/StateSpace/MambaCausality.lean, which proves that extending the input stream preserves earlier outputs; it rests on the scan algebra in NN/MLTheory/Proofs/StateSpace/Scan.lean.

structure Models.DiagonalS4Spec (α : Type) [TorchLean.Storage α] (inputDim stateDim outputDim : ℕ) : Type

Parameters for a diagonal S4-style sequence layer.

• inProj : TorchLean.Tensor α [inputDim, stateDim]

  Input projection from token/features into SSM state channels.

• outProj : TorchLean.Tensor α [stateDim, outputDim]

  Output projection from SSM channels to token/features.

• ssm : Spec.Dynamics.DiagonalSSM α stateDim

  Diagonal recurrent state-space core.

def Models.DiagonalS4Spec.projectInput {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {inputDim stateDim outputDim : ℕ} (m : DiagonalS4Spec α inputDim stateDim outputDim) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [stateDim]

Project an input token into state channels.

def Models.DiagonalS4Spec.projectOutput {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {inputDim stateDim outputDim : ℕ} (m : DiagonalS4Spec α inputDim stateDim outputDim) (h : TorchLean.Tensor α [stateDim]) : TorchLean.Tensor α [outputDim]

Project state channels to output channels.

def Models.DiagonalS4Spec.step {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {inputDim stateDim outputDim : ℕ} (m : DiagonalS4Spec α inputDim stateDim outputDim) (h : TorchLean.Tensor α [stateDim]) (x : TorchLean.Tensor α [inputDim]) : TorchLean.Tensor α [stateDim] × TorchLean.Tensor α [outputDim]

One recurrent S4-style token step, returning (new_state, output).

def Models.DiagonalS4Spec.runArray {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {inputDim stateDim outputDim : ℕ} (m : DiagonalS4Spec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) (xs : Array (TorchLean.Tensor α [inputDim])) : TorchLean.Tensor α [stateDim] × Array (TorchLean.Tensor α [outputDim])

Run an array of tokens through the recurrent layer.

@[simp]

theorem Models.DiagonalS4Spec.runArray_empty {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {inputDim stateDim outputDim : ℕ} (m : DiagonalS4Spec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) : m.runArray h0 #[] = (h0, #[])

An empty token sequence leaves the S4 hidden state untouched and emits nothing.

@[simp]

theorem Models.DiagonalS4Spec.runArray_outputs_size {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {inputDim stateDim outputDim : ℕ} (m : DiagonalS4Spec α inputDim stateDim outputDim) (h0 : TorchLean.Tensor α [stateDim]) (xs : Array (TorchLean.Tensor α [inputDim])) : (m.runArray h0 xs).2.size = xs.size

A recurrent S4 pass emits one output token per input token.