Layer Construction #
This module defines explicit-seed layer builders under TorchLean.nn.Internal. The seeded builders
in NN.API.Seeded allocate initialization seeds from a deterministic stream.
Sequential model type (TorchLean Seq), analogous to PyTorch nn.Sequential.
Instances For
Expose common Seq helpers under TorchLean.nn.
The names mirror the TorchLean runtime layer so users can move between the public API and runtime layer code without learning a second vocabulary.
Lift a single layer definition into a sequential model.
Instances For
All explicit-seed layer constructors live under nn.Internal.*.
The top-level nn.* namespace is reserved for the seeded builder API that allocates
initialization seeds automatically (PyTorch-style ergonomics).
Convert a layer-like value to a sequential model for seq! composition.
- asSequential {σ τ : Spec.Shape} : F σ τ → Sequential σ τ
Instances
Compose layers and sequential models accepted by the seq! syntax.
Instances For
Parameter initialization for an affine layer. none selects Xavier-uniform weights.
- weightInit? : Option Runtime.Autograd.Torch.Init.Scheme
- biasInit : Runtime.Autograd.Torch.Init.Scheme
Instances For
Linear layer on the last axis (prefix-shape preserving).
PyTorch analogue: torch.nn.linear.
See https://pytorch.org/docs/stable/generated/torch.nn.linear.html.
Unlike the runtime TorchLean layer constructor (which is vector-only), this public layer constructor follows PyTorch’s convention:
- if
xhas shape[..., inDim],linear inDim outDimreturns a model of shape[..., outDim].
The leading “prefix” dimensions are treated as a batch (they are flattened to (numel(prefix), inDim),
the affine map is applied once, and the result is reshaped back).
Instances For
Linear layer with Xavier-uniform weights and zero bias.
Instances For
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 seqLen steps using existing TorchLean ops, so it runs on both
CPU and CUDA backends.
PyTorch analogy: torch.nn.RNN(inputSize, hiddenSize, nonlinearity="tanh") with
batch_first=false, specialized to a single batch element.
Instances For
GRU layer (time-major sequence, no batch axis).
This is implemented by unrolling seqLen steps using existing TorchLean ops, so it runs on both
CPU and CUDA backends.
PyTorch analogy: torch.nn.GRU(inputSize, hiddenSize) with batch_first=false, specialized to a
single batch element.
Instances For
Trainable Mamba-style gated diagonal state-space layer.
The layer is time-major and single-batch, matching the simple rnn/gru/lstm constructors:
input (seqLen × inputSize), output (seqLen × hiddenSize). It is unrolled with differentiable
TorchLean ops, so CPU and CUDA training use the same API.
Instances For
LSTM layer (time-major sequence, no batch axis).
This is implemented by unrolling seqLen steps using existing TorchLean ops, so it runs on both
CPU and CUDA backends.
PyTorch analogy: torch.nn.LSTM(inputSize, hiddenSize) with batch_first=false, specialized to a
single batch element.
Instances For
Embedding table initialization configuration (one-hot / token-distribution inputs).
TorchLean-friendly analogue of torch.nn.Embedding in the common setting where token ids are
represented as one-hot vectors (or soft token distributions), so lookup is a matrix multiplication
rather than integer indexing.
- seedW : ℕ
Seed for deterministic embedding-table initialization.
Initialization scheme for the embedding table.
Instances For
Embedding layer for one-hot / token-distribution inputs (no bias).
Input shape: [..., vocab]
Output shape: [..., embedDim]
PyTorch analogue: conceptually nn.Embedding(vocab, embedDim) but applied to one-hot inputs.
Instances For
Learned positional embedding configuration.
This is a trainable parameter tensor of shape (seqLen × embedDim) that is broadcast across the
leading batch dimension and added to the input.
- seedPos : ℕ
Seed for deterministic initialization.
- posInit : Runtime.Autograd.Torch.Init.Scheme
Initialization scheme for the positional embedding table.
Instances For
Add learned positional embeddings to a batched (batch × seqLen × embedDim) tensor.
PyTorch analogue: x + pos[:seqLen] where pos is a parameter table.
Instances For
Add sinusoidal positional encodings to a batched (batch × seqLen × embedDim) tensor.
Implementation:
- precompute
PE : (seqLen × embedDim)at initialization time (stored as a non-trainable buffer), - broadcast it across the leading
batchaxis and add to the input.
Instances For
Apply RoPE to a batched multi-head tensor (batch × numHeads × seqLen × headDim).
This matches the standard identity:
$$ \operatorname{rope}(x) = x \odot \cos + \operatorname{rotatePairs}(x) \odot \sin $$
where cos/sin depend only on (pos, dim) and broadcast across (batch, numHeads).
Notes:
- This layer is differentiable (gradients flow through the rotation), but it has no trainable
parameters; the precomputed
cos/sintables are stored as non-trainable buffers. - The pure spec version is in
NN.Spec.Layers.PositionalEncoding(Spec.rope_apply_heads_spec).
Instances For
Reduce-sum to a scalar. PyTorch analogue: torch.sum.
Instances For
Flatten any tensor into a 1D vector of length size s. PyTorch analogue: torch.flatten.
Instances For
Flatten a batched tensor N × σ into a matrix N × (size σ).
PyTorch analogue: torch.flatten(x, start_dim=1).
Instances For
Dropout layer (active in train mode, identity in eval mode).
PyTorch analogue: torch.nn.Dropout.
Instances For
Convenience block: Flatten -> Linear.
This is common for "image to classifier head" models.