TorchLean NN: Linear and Recurrent Layers

Convenience constructors (layers)

def Runtime.Autograd.Model.Layers.Internal.writeLeading {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {rows : ℕ} {tail : Spec.Shape} (base : Ref (tail.prependDim rows)) (value : Ref tail) (index : Fin rows) : m (Ref (tail.prependDim rows))

Write one leading-axis slice through the general scatterAdd operation.

def Runtime.Autograd.Model.Layers.Internal.requirePositive (kind field : String) (value : ℕ) : Except String Unit

Validate one positive architectural dimension.

def Runtime.Autograd.Model.Layers.Internal.validateRecurrentDimensions (kind : String) (sequenceLength inputWidth hiddenWidth : ℕ) : Except String Unit

Validate the dimensions shared by recurrent layers.

def Runtime.Autograd.Model.Layers.linear (inputWidth outputWidth : ℕ) (weightSeed : ℕ := 0) : Layer [inputWidth] [outputWidth]

Fully-connected affine layer on vectors: $y=Wx+b$.

Parameters:

• W : (outputWidth × inputWidth) initialized with Xavier initialization,
• b : (outputWidth) initialized to zeros.

PyTorch analogy: torch.nn.Linear(inputWidth, outputWidth).

def Runtime.Autograd.Model.Layers.rnn (sequenceLength inputWidth hiddenWidth : ℕ) (weightSeed : ℕ := 0) : Layer [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

Vanilla RNN layer (time-major sequence, no batch axis).

Semantics: $$ h_t=\tanh\!\left(W[x_t;h_{t-1}]+b\right), \qquad h_{-1}=0. $$

This is implemented by unrolling a fixed number of steps (sequenceLength) using existing TorchLean ops, so it works on both CPU and CUDA backends.

PyTorch analogy: torch.nn.RNN(inputWidth, hiddenWidth, nonlinearity="tanh") with batch_first=false, specialized to a single batch element. Docs: https://docs.pytorch.org/docs/stable/generated/torch.nn.RNN.html

def Runtime.Autograd.Model.Layers.gru (sequenceLength inputWidth hiddenWidth : ℕ) (resetWeightSeed updateWeightSeed candidateWeightSeed : ℕ := 0) : Layer [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

GRU layer (time-major sequence, no batch axis).

The Cho-style reset-before recurrence starts with $h_{-1}=0$ and computes

$$ \begin{aligned} r_t &= \operatorname{sigmoid}(W_r[x_t;h_{t-1}]+b_r),\\ z_t &= \operatorname{sigmoid}(W_z[x_t;h_{t-1}]+b_z),\\ n_t &= \tanh\!\left(W_{nx}x_t+W_{nh}(r_t\odot h_{t-1})+b_n\right),\\ h_t &= (1-z_t)\odot n_t+z_t\odot h_{t-1}. \end{aligned} $$

The candidate matrix is $W_n=[W_{nx}\;W_{nh}]$: the reset multiplies the previous hidden state before the hidden columns of this matrix are applied. PyTorch uses reset-after instead:

$$ n_t = \tanh\!\left(W_{nx}x_t+b_{nx}+r_t\odot(W_{nh}h_{t-1}+b_{nh})\right). $$

In that equation the reset multiplies the recurrent affine output, including the recurrent candidate bias $b_{nh}$. A general matrix does not commute with elementwise reset multiplication, and $r_t\odot b_{nh}$ is not a constant bias. Concatenating or reordering checkpoint tensors and adding their biases therefore cannot generally turn a PyTorch GRU into this cell.

The six trainable tensors are stored as wReset, bReset, wUpdate, bUpdate, wNew, bNew. Every weight has shape [hiddenWidth, inputWidth + hiddenWidth], with input columns first and hidden columns second; every bias has shape [hiddenWidth]. PyTorch packs reset, update and candidate rows into separate input and recurrent matrices, with a separate bias vector for each matrix. Its reset and update bias pairs may be summed to reproduce those gates' forward equations. The candidate bias and reset placement require the different recurrence shown above. Even where two biases can be merged for forward evaluation, training one merged bias differs from updating two independent bias parameters.

Weights use Xavier uniform initialization with the three supplied seeds; biases start at zero. Each forward call unrolls [sequenceLength, inputWidth] into [sequenceLength, hiddenWidth] and starts from a fresh zero hidden state. The result contains every hidden state, including the last one as its final row. The layer accepts no initial hidden state, returns no separate final state, and carries no hidden state between calls. The public batched wrapper applies this same core independently to each sequence with shared parameters.

def Runtime.Autograd.Model.Layers.gruResetAfterCell {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {inputWidth hiddenWidth : ℕ} (inputWeight : Ref [3 * hiddenWidth, inputWidth]) (hiddenWeight : Ref [3 * hiddenWidth, hiddenWidth]) (inputBias hiddenBias : Ref [3 * hiddenWidth]) (input : Ref [inputWidth]) (previous : Ref [hiddenWidth]) : m (Ref [hiddenWidth])

One differentiable reset-after GRU step with PyTorch's four packed parameter tensors.

The input and hidden affine maps each produce reset, update, and candidate blocks. We apply the reset to the whole recurrent candidate block, including its bias. All four parameter tensors and the previous hidden state remain ordinary autograd references, so the same program supports copied-parameter evaluation, reverse-mode gradients, and differentiation through several steps.

def Runtime.Autograd.Model.Layers.gruResetAfter (sequenceLength inputWidth hiddenWidth : ℕ) (inputWeightSeed hiddenWeightSeed : ℕ := 0) : Layer [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

GRU sequence layer using the reset-after convention.

State order is weight_ih, weight_hh, bias_ih, bias_hh, with reset/update/candidate rows in each tensor. This is also the order accepted by Spec.GRUResetAfterSpec.ofPyTorch; copied tensors need no gate permutation or bias merging. Each call starts at zero and returns every hidden state. Use gruResetAfterCell when the initial state is supplied by another part of the model.

def Runtime.Autograd.Model.Layers.gruFromPyTorch (sequenceLength : ℕ) {inputWidth hiddenWidth : ℕ} (parameters : Spec.GRUResetAfterSpec Float inputWidth hiddenWidth) : Layer [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

Build a reset-after sequence layer from a single PyTorch cell's copied parameters.

The initialization override uses the supplied tensors directly. There is no random initializer left to replace them at runtime, and the two bias vectors remain separate trainable state slots.

def Runtime.Autograd.Model.Layers.lstm (sequenceLength inputWidth hiddenWidth : ℕ) (forgetWeightSeed inputWeightSeed candidateWeightSeed outputWeightSeed : ℕ := 0) : Layer [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

LSTM layer (time-major sequence, no batch axis).

This is an unrolled LSTM using the standard four gates, with $(h_{-1},c_{-1})=(0,0)$.

PyTorch analogy: torch.nn.LSTM(inputWidth, hiddenWidth) with batch_first=false, specialized to a single batch element. Docs: https://docs.pytorch.org/docs/stable/generated/torch.nn.LSTM.html