Differentiable diagonal selective scans

Each step computes h[t] = A[t] * h[t-1] + B[t] * X[t], with elementwise products. The result contains every state after its corresponding input. The generic interpreter records the recurrence using ordinary differentiable operations; the native interpreter records one scan node and propagates cotangents to coefficients, inputs, and the initial state.

The final row can be passed as the initial state of another call without cutting its gradient. An empty sequence returns an empty result and has no dependence on the initial state.

def Runtime.Autograd.Model.F.SelectiveScan.variableReference {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {state seqLen : ℕ} : RefTy m α [seqLen, state] → RefTy m α [seqLen, state] → RefTy m α [seqLen, state] → RefTy m α [state] → m (RefTy m α [seqLen, state])

Reference recurrence with token-dependent coefficients and an explicit initial state.

def Runtime.Autograd.Model.F.selectiveScanDiagVar {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {seqLen state : ℕ} (a b x : RefTy m α [seqLen, state]) (initial : RefTy m α [state]) : m (RefTy m α [seqLen, state])

Diagonal selective scan with independent coefficients for every token.

All three sequence inputs have shape [seqLen, state]. Gradients flow through A, B, X, and initial, including when the coefficients themselves were computed from the input tokens.

def Runtime.Autograd.Model.F.selectiveScanDiag {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {seqLen state : ℕ} (a b : RefTy m α [state]) (x : RefTy m α [seqLen, state]) (initial : RefTy m α [state]) : m (RefTy m α [seqLen, state])

Diagonal scan with coefficient vectors shared across time.

Reverse mode sums the contributions to each shared coefficient over all steps. The native kernel performs that accumulation directly; the reference route gets it from broadcasting.