Crate axis 

Experimental training with named axes on a stream-ordered cuTile GPU backend.

Muon is adapted from Keller Jordan’s reference implementation. The pinned source revision and MIT notice are packaged in THIRD_PARTY.md.

Modules

architectures

Reference architectures assembled only from public library modules: composition witnesses as well as convenience constructors. One file per family under src/architectures/. Equivalence to each architecture’s reference implementation is established by tests in src/library/src/tests.rs, never asserted here.

prelude

Structs

Adam

Device-resident Adam with bias correction and one state pair per ParamId.

AdamW

Adam with decoupled weight decay; moments and updates remain device-resident.

AdaptiveAvgPool1d

Named-axis adaptive average pooling over one spatial axis: reduces it to a fixed target extent using PyTorch’s own per-axis bin formula. See Tensor::adaptive_avg_pool3d for the exact bin math and its uneven-division behavior.

AdaptiveAvgPool2d

Named-axis adaptive average pooling over two spatial axes; see AdaptiveAvgPool1d and Tensor::adaptive_avg_pool3d for the shared bin contract.

AdaptiveAvgPool3d

Named-axis adaptive average pooling over three spatial axes, wrapping Tensor::adaptive_avg_pool3d as a stateless Module; see AdaptiveAvgPool1d for the shared contract.

AdaptiveMaxPool1d

Named-axis adaptive max pooling over one spatial axis: reduces it to a fixed target extent using the same bin formula as AdaptiveAvgPool1d, taking each bin’s maximum instead of its mean. See Tensor::adaptive_max_pool3d for the exact composition and its tie-break.

AdaptiveMaxPool2d

Named-axis adaptive max pooling over two spatial axes; see AdaptiveMaxPool1d and Tensor::adaptive_max_pool3d for the shared contract.

AdaptiveMaxPool3d

Named-axis adaptive max pooling over three spatial axes, wrapping Tensor::adaptive_max_pool3d as a stateless Module; see AdaptiveMaxPool1d for the shared contract.

AdditionDataset

An inexhaustible deterministic stream of fresh addition problems. IDs are unique draw counters; values can still coincide after f32 projection.

AdditionSample

AlphaDropout

SELU-preserving alpha dropout: PyTorch’s nn.AlphaDropout. Ordinary inverted dropout resets dropped units to zero, which is off-distribution for a SELU network whose self-normalizing property assumes zero mean and unit variance; alpha dropout instead resets dropped units to SELU’s own saturation value and applies an affine correction that keeps the first two moments fixed. Formula (bit-for-bit PyTorch’s aten/src/ATen/native/Dropout.cpp _alpha_dropout_impl, alpha_dropout=true): with alpha = 1.7580993408473766 (SELU’s -alpha_selu * scale_selu, negated) and a = 1 / sqrt((alpha^2 * p + 1) * (1 - p)), each element keeps with probability 1 - p (mask keep = draw >= p, the same Tensor::uniform_device convention Dropout uses) and the output is a * input * keep + alpha * a * keep + alpha * a * (p - 1) – algebraically a * (input + alpha) where kept, and the constant alpha * a * (p - 1) where dropped, with no dedicated “add a scalar constant” primitive needed: the three terms above compose from mul, scale and add alone, and the constant term is built with Tensor::from_slice exactly like normalization::Affine’s own constant parameters. forward (evaluation) is the identity. forward_training consumes exactly one pass.next_seed() draw regardless of p; p == 0 is the identity and p == 1 is exact zeros with exactly zero gradient (PyTorch’s own _dropout_impl special-cases p == 1 to input * 0 before the affine constants – which are undefined at p == 1 – are ever computed, for every dropout variant including this one).

AvgPool1d

Named-channel 1D average pooling: sum(window) / divisor, independently per channel, matching PyTorch’s nn.AvgPool1d and all of its options. The divisor is divisor_override when set; otherwise, with count_include_pad=True (the default), the window’s extent clipped to the padded input – the full kernel extent except for a ceil_mode window overhanging the right padding; with count_include_pad=False, the count of real (non-padding) positions. The window count is floor((L + 2p - k) / s) + 1, or with ceil_mode=True ceil((L + 2p - k) / s) + 1 minus one when that last window would start at or past L + p (PyTorch’s own rule: every window starts inside the input or its left padding). Stride defaults to the kernel extent and padding to 0; padding must be at most half the kernel extent, the same constraint MaxPool1d enforces.

AvgPool2d

Named-channel 2D average pooling; see AvgPool1d for the shared ceil_mode, count_include_pad and divisor_override contract, applied per axis over two spatial axes as MaxPool2d is; a window’s divisor is the product of its per-axis counts.

AvgPool3d

Named-channel 3D average pooling; see AvgPool1d for the shared contract, applied over three spatial axes as MaxPool3d is.

Axis

Identity is independent of both the diagnostic name and tensor-local extent.

Batch

BatchNorm

Batch normalization over one explicitly named feature axis, PyTorch’s BatchNorm1d/BatchNorm2d/BatchNorm3d in a single named-axis module: every OTHER axis present at call time – batch and any spatial axes – normalizes together, whatever it is named. Defaults to epsilon 1e-5, momentum Some(0.1), a learnable per-feature affine scale and bias, and track_running_stats = true, matching PyTorch’s own defaults.

Bilinear

Learned bilinear form over two named axes, y = x1^T A x2 + b, PyTorch’s nn.Bilinear. Composed from two Tensor::contract calls – x1 contracted against weight over in1, then that intermediate contracted against x2 over in2 – the same contraction Linear already uses, applied twice, so it needs no dedicated kernel. Any axis x1 and x2 share beyond in1/in2 (such as a common batch axis) aligns automatically through that shared contraction, exactly like Linear’s own unrelated axes; an axis unique to either input passes through untouched.

CELU

Elementwise CELU with an explicit, finite, nonzero alpha – PyTorch’s own documented domain, “valid for alpha != 0”, negative alpha included. See Tensor::celu for the exact forward/backward split at x == 0 and why the two-branch composition already matches PyTorch’s max(0, x) + min(0, ...) definition for any nonzero alpha.

CategoricalAccuracy

Exact categorical classification counts. Counts compose across evaluation batches.

CentralDifference

A centered finite-difference stencil evaluated at a positive coordinate step.

ChannelDropout

Channel (feature-map) dropout: PyTorch’s nn.Dropout1d, nn.Dropout2d and nn.Dropout3d are one type here, since Axis names axes instead of positionally counting spatial dimensions – exactly how ZeroPad already collapses PyTorch’s ZeroPad1d/2d/3d into one axis-list constructor: the “1d”/“2d”/“3d” split is purely how many axes are named spatial, never a different type or contract. Drops one mask value per combination of channel and every axis NOT named spatial (so every batch-like axis varies independently too, matching PyTorch’s make_feature_noise mask shape [N, C, 1, 1, ...]), broadcasting that one decision over every spatial axis: Conv2d’s (batch, channel, height, width) output passed through ChannelDropout::new(p, channel, [height, width]) zeroes whole feature maps together, the way dropping isolated pixels would not regularize a convolutional feature. Otherwise identical to Dropout: forward is the identity; forward_training consumes exactly one pass.next_seed() draw regardless of p, feeds it to one Tensor::uniform_device call over the non-spatial axes alone, keeps each (channel, batch, …) combination whose draw is >= p, and rescales kept elements by 1 / (1 - p); p == 0 is the identity and p == 1 is exact zeros with exactly zero gradient, the same edge cases Dropout documents.

ChannelShuffle

torch.nn.ChannelShuffle(groups): split channel (extent C) into [group (groups), within (C / groups)] – PyTorch’s actual internal reshape order, (N, groups, C/groups, *), which is the reverse of its own prose (“divides … into g groups as (N, C/g, g, *)”), verified against a real channel_shuffle run rather than trusted from the text – then merge back in swapped order [within, group], reproducing its documented transpose-and-flatten exactly ([ch0, ch1, ch2, ch3] at groups=2 becomes [ch0, ch2, ch1, ch3], matching the PyTorch docs’ own example; the two group/within orderings only coincide when groups == C / groups, which is why that example alone cannot distinguish them). Every other axis is untouched, so this works on any (N, C, *) shape without naming the trailing axes. No dedicated kernel or backward rule: split and merge are both already differentiable.

CircularPad

Circular (wrap-around) padding, covering CircularPad1d/2d/3d. Matches PyTorch’s own constraint that each side’s padding be at most the axis’s extent (a full wrap is allowed, unlike reflection’s strict <); see circular_pad_axis for the composition and its exact gradient.

ConstantPad

Padding at an arbitrary constant value, covering ConstantPad1d/2d/3d. ZeroPad’s own zero-padded tensor already has the correct interior (pad_zeros never touches it) and exact zeros in the border; a second, independent zero-padding of a constant ones mask (Tensor::zeros(..).ge(0.0), which carries no gradient edge of its own) turns into an interior/border indicator, and its logical_not – exactly zero inside, exactly one in the border – scaled by value and added supplies the constant. Backward is therefore identical to ZeroPad’s: the border term is a detached constant and contributes no gradient.

Conv1d

Named-channel 1D cross-correlation. Reuses Conv2d’s exact geometry, weight layout, and xorshift initialization over a synthetic unit spatial axis appended for the call: a kernel of 1, stride of 1, and padding of 0 on that axis is an exact no-op (it always contracts a single, always-in-bounds position), and Tensor::broadcast_to/Tensor::select add and remove an extent-1 axis with an exact identity gradient, so no separate rank-1 backend path exists. Stride and padding default to 1 and 0; groups default to 1, matching Conv2d.

Conv2d

Named-channel 2D cross-correlation followed by bias addition.

Conv3d

Named-channel 3D cross-correlation followed by bias addition.

ConvTranspose1d

Named-channel 1D transposed convolution; see ConvTranspose2d for the shared contract. Reuses ConvTranspose2d’s exact geometry and initialization over a synthetic unit spatial axis, exactly as Conv1d reuses Conv2d.

ConvTranspose2d

Named-channel 2D transposed convolution (PyTorch’s nn.ConvTranspose2d): stride, padding, output_padding and groups are supported; dilation is pinned at 1, matching Conv2d. Stride defaults to [1, 1], padding and output_padding to [0, 0], and groups to 1.

ConvTranspose3d

Named-channel 3D transposed convolution; see ConvTranspose2d for the shared contract. The three spatial axes conventionally represent depth, height, and width, matching Conv3d.

DataLoader

Device

Dim

Disjointness

Exact separation across two or more declared populations.

DisjointnessReceipt

Dropout

Parameter-free inverted dropout (PyTorch’s nn.Dropout). forward (evaluation) is the identity: it returns the same tensor unchanged, never a copy scaled by 1. forward_training consumes exactly one pass.next_seed() draw per call, regardless of p, so a downstream random consumer’s draw sequence does not shift when p changes; that seed only feeds an actual Tensor::uniform_device(shape, seed, device) draw when 0 < p < 1. The mask is generated entirely on the device – never a host-drawn tensor uploaded over PCIe every step – from a counter-based hash of (seed, element index) (see Tensor::uniform_device, backend::kernels::uniform_device); it keeps each element whose draw is >= p, zeroes the rest, and rescales the kept elements by 1 / (1 - p) (PyTorch’s inverted dropout), so the output’s expectation matches the input. p == 0 returns the input unchanged; p == 1 returns exact zeros with exactly zero gradient (PyTorch’s own p == 1 behavior), since 1 / (1 - p) is undefined there. Elsewhere the gradient is mask / (1 - p) by construction – a constant {0, 1} mask times a scalar – so no dedicated backward rule is needed.

ELU

Elementwise ELU with an explicit, finite, nonzero alpha (PyTorch’s own domain – negative alpha included). See Tensor::elu for the exact forward/backward split at x == 0 and the sign-algebra argument for why negative alpha is safe here.

Embedding

Look up one learned feature vector per vocabulary entry.

EmbeddingBag

One learned feature vector per vocabulary entry, pooled per bag: PyTorch’s nn.EmbeddingBag. Unlike Embedding’s dense one-hot contraction, the input is a host-side flat row-index array together with offsets (bag b covers positions offsets[b]..offsets[b + 1], or ..index.len() for the last bag – exactly torch.nn.EmbeddingBag.forward’s own input/ offsets pair), so it scales to a large table the same way Tensor::gather itself does. Does not implement Module, which only threads a single tensor input: the index and offsets are host-side integer arrays, not tensors, exactly as Tensor::gather’s own index and Tensor::scatter_add’s own index/bucket already are.

EmpiricalMonotonicity

A sampled check, never a global proof about the represented function.

EmpiricalMonotonicityReceipt

EmpiricalResidual

A sampled residual check, never a proof that the represented law holds.

EmpiricalResidualReceipt

ExactGELU

Parameter-free erf-form GELU for checkpoints trained with exact GELU, which is PyTorch’s default nn.GELU().

FeatureAlphaDropout

Channel-wise alpha dropout: PyTorch’s nn.FeatureAlphaDropout, AlphaDropout’s exact affine formula applied to a ChannelDropout-shaped mask (one draw per (channel, batch, …) combination, broadcast over the named spatial axes) instead of one draw per element. See ChannelDropout for the axis contract and AlphaDropout for the formula; this module is their composition, matching PyTorch’s own _feature_alpha_dropout (feature_dropout=true, alpha_dropout=true in the same _dropout_impl).

FinitePasses

Exact accounting for a declared number of complete finite-corpus passes.

FinitePassesLoader

Batches a finite population for exactly N independently shuffled complete passes.

FinitePassesReceipt

Flatten

Merge a run of named axes into one, PyTorch’s nn.Flatten expressed over identities instead of a positional start_dim/end_dim range: the caller names the exact axes to fold, in the physical order they should flatten in, and the single output axis that replaces them. A thin Module wrapper over Tensor::merge; see its own doc for the exact selected-axis-order and physical-layout contract this reuses verbatim, including which position the merged axis is inserted at.

Fold

Named-axis col2im (PyTorch’s nn.Fold), the adjoint of Unfold: combines an array of sliding local kernel-sized blocks (patch in place of a channel axis) by summing overlapping contributions into output_size-shaped spatial axes, reconstructing channels (inferred from the input’s patch extent divided by the kernel volume). Rejects an input whose extents are inconsistent with output_size/kernel/stride/padding before any device call. Stride defaults to [1, 1] and padding to [0, 0].

GELU

Parameter-free tanh-form GELU. PyTorch’s default nn.GELU() is the erf form, ExactGELU; a faithful port of an un-annotated PyTorch model uses that one.

GLU

Gated Linear Unit over an explicit named axis, halved by the gate. See Tensor::glu.

GroupNorm

Group normalization over channel groups and explicitly named sample axes.

Gru

A stacked, optionally bidirectional, eager GRU over one named time axis. See Lstm’s doc comment for the eager-transition contract and RecurrentConfig for num_layers/bidirectional/bias. The default configuration (one layer, unidirectional, bias on) is bit-exact with the single-layer GRU this type used to be.

GruCell

One PyTorch-form GRU transition with explicit hidden state.

GruRun

Sequence output and the per-layer, per-direction terminal hidden states from one Gru run, ordered as LstmRun::state documents.

Hardshrink

Parameter-free Hardshrink with an explicit lambd. See Tensor::hardshrink for the exact forward and gradient-at-the-band semantics.

Hardsigmoid

Parameter-free Hardsigmoid. See Tensor::hardsigmoid for the exact forward and gradient-at-the-kinks semantics.

Hardswish

Parameter-free Hardswish. See Tensor::hardswish for the exact forward and gradient-at-the-kinks semantics, including its asymmetric boundary convention.

Hardtanh

Parameter-free Hardtanh with explicit min_val/max_val bounds. See Tensor::hardtanh for the exact forward and gradient-at-the-kinks semantics.

Identity

Parameter-free passthrough, PyTorch’s nn.Identity: returns its input unchanged.

IdentityScheme

Names the equivalence relation used to decide whether two observations are the same scientific unit.

Idr

Exact observed-ID accounting. Callers choose a stable ID whose equality means reuse.

IdrLimits

Declared operational limits for an infinite-data approximation.

IdrReceipt

InMemoryDataset

InstanceNorm

Per-channel instance normalization over explicitly named sample axes.

JitCacheStats

Cumulative cuTile JIT disk-cache hit/miss counts for this process, zero for both fields before any CUDA device has compiled a kernel, or when the cache is disabled (AXIS_JIT_CACHE=off) or could not be enabled. See Device::cuda.

LPPool1d

Named-channel 1D power-average pooling: (sum(x^p))^(1/p) over each kernel window, independently per channel, matching PyTorch’s nn.LPPool1d (f(X) = (sum_{x in X} x^p)^(1/p); at p = 1 this is exactly sum pooling, and there is no division by the window size). p (PyTorch’s norm_type) is any finite positive real. An integer p uses repeated products; a real p uses exp(p ln x) with 0^p = 0, so a negative input under a non-integer p is NaN, as x.pow(p) is in PyTorch. One deliberate difference remains: an odd integer p whose window sum is negative returns the real signed root sign(s) |s|^(1/p), where PyTorch’s pow(1/p) returns NaN. With ceil_mode=True, windows are counted as in AvgPool1d and a window clipped at the input edge is (sum * k / clipped)^(1/p), exactly PyTorch’s avg_pool(x^p) * k composition. PyTorch’s LPPool has no padding parameter, so none is exposed here either. Stride defaults to the kernel extent, matching every other pooling family in this file.

LPPool2d

Named-channel 2D power-average pooling; see LPPool1d for the shared contract, applied over two spatial axes as MaxPool2d is.

LPPool3d

Named-channel 3D power-average pooling; see LPPool1d for the shared contract, applied over three spatial axes as MaxPool3d is.

LawIdentity

Stable identity for the scientific law whose residual is being evaluated.

LayerNorm

Population layer normalization over one or more declared named axes.

LeakyReLU

Parameter-free leaky ReLU with an explicit negative-region slope. The derivative at exactly zero is the negative slope, matching PyTorch.

LearningLimits

LearningObservation

LearningProgress

Compares one declared metric on one declared population over an ordered budget.

LearningProgressReceipt

Linear

Contract input, introduce output, and preserve all unrelated axes. Cloning a built Linear explicitly ties its parameters. Carries a learned [output] bias by default; .bias(false) before build omits it entirely, so named_parameters then has only weight and no zero-filled bias tensor is ever allocated.

LinearCrossEntropyOptions

Options for Tensor::linear_cross_entropy_with_logits, mirroring torch.nn.functional.linear_cross_entropy(input, linear_weight, target, *, linear_bias=None, weight=None, reduction='mean', ignore_index=None, label_smoothing=0.0, options=None) (PyTorch 2.14, added to fuse a final Linear projection with cross_entropy so the full [*, n_classes] logits tensor is never materialized at large vocabularies). Only label_smoothing is exposed here: reduction is left to the caller via Tensor::mean/Tensor::sum on the class-reduced result, matching every other loss in this file (multilabel_soft_margin_loss, binary_cross_entropy_with_logits); per-class weight and ignore_index are not implemented (this composition uses soft/one-hot probability targets, matching categorical_cross_entropy_with_logits, not sparse class-index targets with a sentinel to ignore); and PyTorch’s options= chunked fast path has no analogue (see linear_cross_entropy_with_logits’s own doc comment for why).

LocalResponseNorm

Response normalization across a fixed-size sliding window of a named channel axis, PyTorch’s nn.LocalResponseNorm: each element divides by (k + alpha / size * sum(a_c'^2))^beta, where the sum runs over a window of size channels around c, padded with implicit zeros at the channel-axis boundary (size / 2 channels before c, (size - 1) / 2 after, PyTorch’s own asymmetric split for an even size). No learned parameters, matching PyTorch’s own stateless module.

LogSigmoid

Parameter-free numerically stable log-sigmoid. See Tensor::log_sigmoid.

LogSoftmax

Numerically stable log-softmax over an explicit named axis. See Tensor::log_softmax.

Lstm

A stacked, optionally bidirectional, eager LSTM over one named time axis.

LstmCell

One standard IFGO LSTM transition with explicit hidden and cell state.

LstmRun

Sequence output and the per-layer, per-direction terminal recurrent states from one Lstm run. state[layer * directions + direction] (direction 0 = forward, 1 = backward when bidirectional), matching PyTorch’s (num_layers * num_directions, batch, hidden) stacking order.

LstmState

Explicit recurrent state. Both tensors contain the same named stream axes and the LSTM’s hidden axis.

MaxPool1d

Named-channel 1D max pooling; see MaxPool2d for the shared contract, applied along one spatial axis.

MaxPool2d

Named-channel 2D max pooling: reduce each kernel window to its maximum, independently per channel. Stride defaults to the kernel extent and padding to [0, 0], matching PyTorch’s nn.MaxPool2d. Spatial arrays are ordered as the supplied named axes, preserved (at their new, pooled extent) in their original position; the channel axis is appended, exactly as Conv2d appends its output-channel axis. Padding must be at most half the kernel extent in each dimension (PyTorch’s own MaxPool constraint); a padded position can never win the maximum. Ties route the gradient to the first logical coordinate along each spatial axis, independently of physical layout, matching Tensor::min.

MaxPool3d

Named-channel 3D max pooling; see MaxPool2d for the shared contract. The three spatial axes conventionally represent depth, height, and width, and its kernel/stride/padding arrays use that same supplied order, matching Conv3d.

MaxUnpool1d

Named-axis 1D max unpooling, PyTorch’s nn.MaxUnpool1d: the partial inverse of MaxPool1d, scattering each input value to the position its index names in a zero output and leaving every other position zero. indices are MaxPool1d::forward_with_indices’s, one per input element in its logical order; see Tensor::max_unpool3d for duplicates and the gradient (a gather of the upstream gradient by the same indices). Stride defaults to the kernel extent and padding to 0; the output extent defaults to (in - 1) * stride - 2 * padding + kernel, and an explicit one (to undo a pool whose floor division dropped a remainder) must lie strictly within stride of that default, as PyTorch requires. Not a Module: like PyTorch’s, its forward takes the indices as a second input.

MaxUnpool2d

Named-axis 2D max unpooling, PyTorch’s nn.MaxUnpool2d; see MaxUnpool1d for the shared contract. Its indices are MaxPool2d::forward_with_indices’s, row-major offsets into the output’s spatial plane in the supplied spatial order.

MaxUnpool3d

Named-axis 3D max unpooling, PyTorch’s nn.MaxUnpool3d; see MaxUnpool1d for the shared contract. Its indices are MaxPool3d::forward_with_indices’s.

Mish

Parameter-free Mish, x * tanh(softplus(x)). See Tensor::mish.

ModuleDict

Named collection of modules with no forward of its own, insertion-ordered like PyTorch’s OrderedDict-backed nn.ModuleDict. See ModuleList’s own doc for the shared no-forward/aggregation contract; slots here are addressable as "{key}.{name}". Duplicate keys are rejected before any module is built.

ModuleList

Ordered collection of modules with no forward of its own, PyTorch’s nn.ModuleList: PyTorch’s own nn.Module.forward on it raises NotImplementedError, since only a subclass that actually threads data through its children – such as Sequential – has one, so output_shape/build/forward below each return an explicit error instead of guessing an order. Build and run a held module directly through ModuleList::get/ModuleList::get_mut/ModuleList::iter. named_parameters still aggregates every held module’s own parameters under Sequential’s own slot-path convention, "{index}.{name}" (e.g. "0.weight"), so they stay addressable through model.parameter(path) even though nothing here threads a single input through them in order. ModuleDict, ParameterList, and ParameterDict share this exact contract, keyed or indexed differently.

MonotonicityLimits

MultiheadAttention

Scaled dot-product multi-head attention, torch.nn.MultiheadAttention at its own default dropout=0.0: Attention(Q, K, V) = softmax(QK^T / sqrt(head_feature)) V, computed independently per head and concatenated before one output projection (PyTorch’s forward, batch_first-agnostic since Axis has no positional batch axis to begin with).

Muon

Momentum updates orthogonalized by five Newton-Schulz iterations.

MuonMatrix

A parameter selected for Muon together with its matrix orientation.

MuonWithAuxAdamW

Explicitly routes selected rank-2 parameters to Muon and the exact remainder to AdamW. No shape-based selection occurs.

PReLU

Parametric ReLU with a learnable weight, either shared across every element (PReLU::shared, PyTorch’s default num_parameters=1) or one per entry of an explicit named channel axis (PReLU::channel, PyTorch’s num_parameters=C); both start at PyTorch’s default init, 0.25. See Tensor::prelu.

ParamId

Parameter

Cloning explicitly shares a parameter, including its accumulated gradient.

ParameterDict

Named collection of standalone parameters, not owned by any module: PyTorch’s nn.ParameterDict. See ModuleList’s own doc for the shared no-forward/aggregation contract; slots here are addressable by key. Duplicate keys are rejected immediately.

ParameterList

Ordered collection of standalone parameters, not owned by any module: PyTorch’s nn.ParameterList. See ModuleList’s own doc for the shared no-forward/aggregation contract; slots here are addressable by index, "{index}".

PixelShuffle

Sub-pixel upscaling: rearrange channel (extent C * factor^2) into factor blocks along each of the two named spatial axes, expanding each by factor and shrinking channel to C. PyTorch’s own pixel_shuffle decomposition – reshape the channel axis to (C, factor, factor), permute to interleave each factor axis with its spatial axis, then flatten – is exactly Tensor::split (channel into [out_channel, row_factor, col_factor], outermost first) followed by two Tensor::merge calls ([height, row_factor], [width, col_factor], spatial axis outermost so it varies slower than its own sub-pixel offset). No dedicated kernel or backward rule: split/merge are both already differentiable.

PixelUnshuffle

The exact inverse of PixelShuffle: shrink channel’s two named spatial axes by factor each and grow channel by factor^2. Composed as Tensor::split on each spatial axis ([outer, sub_pixel], spatial axis outermost, undoing PixelShuffle’s merge order exactly) followed by one Tensor::merge of [channel, row_factor, col_factor] into the new channel axis (undoing PixelShuffle’s split order exactly).

PopulationLinear

A family of independent linear maps sharing one explicit population axis. Inputs without the population axis are broadcast to every member; inputs already carrying it keep memberwise independence through later layers.

PopulationReceipt

PopulationSpec

PositionEmbedding

Add one learned vector per position: a [position, feature] table whose extents are read from the input at build and broadcast over every other axis. The output shape equals the input shape. Entries start uniform in [-0.02, 0.02).

PrefixDropout

Prefix (nested) dropout (Axis issue #90; no PyTorch class – docs/direction/next.md‘s tracked next implementation): truncates a tensor to its first w positions along one named ordered axis, with w drawn independently per training example, so one trained model yields a whole width sweep at evaluation time (the consumers’ own framing: learning/bae’s ordered sign code and vision/image-encode’s ordered latent/site lists, both cited on the issue). Every axis of the input other than axis itself is an “example” axis: each combination of their coordinates draws its own independent width, exactly like ChannelDropout’s non-spatial axes. forward (evaluation) is the identity – the full, untruncated tensor, matching “evaluation keeps everything.” forward_training consumes exactly one pass.next_seed() draw, feeding one Tensor::uniform_device call over the example axes alone (so the draw is per example, not per element): w = draw * (extent(axis) + 1), a real value in [0, extent + 1), compared against a constant 0 .. extent position tensor along axis (Tensor::broadcast_to combines the two disjoint axis sets – the per-example width and the per-position index – onto one shared shape first, exactly the torch.cdist-style outer-combination idiom broadcast_to’s own doc comment describes). Position i survives exactly when i < w: since w is real-valued and i is an integer, this keeps floor(w) + 1 positions in the generic case (w essentially never lands exactly on an integer), so the kept prefix length is uniform over {1, ..., extent} (extent itself when w lands in its last unit interval; 0 only at the zero-probability exact draw w == 0) – the discrete-uniform-width contract nested.py’s sample_k and stage_a.py’s arange(width) < m both draw, restated without a host round trip. No rescaling is applied: unlike Dropout’s inverted scaling, this is a hard truncation mask with expectation <= 1, not an expectation-preserving Bernoulli mask, matching arange(K) < m in both cited consumers exactly. The kept gradient passes through unchanged and the dropped gradient is exactly zero, the same constant-mask construction Dropout already uses.

RReLU

PyTorch’s nn.RReLU (randomized leaky ReLU, Xu et al. “Empirical Evaluation of Rectified Activations in Convolutional Network”). Defaults lower = 1 / 8, upper = 1 / 3, PyTorch’s own defaults. forward (evaluation) is Tensor::leaky_relu at the fixed slope (lower + upper) / 2, PyTorch’s documented evaluation behavior. forward_training draws a fresh independent slope per element, uniform in [lower, upper), from one pass.next_seed() draw fed to one Tensor::uniform_device call scaled into that range (lower + (upper - lower) * draw, the constant lower shift built with Tensor::from_slice exactly like AlphaDropout’s constant term); every element keeps input unchanged where input > 0 and multiplies by its own per-element slope elsewhere – 0 itself takes the scaled branch, matching PyTorch’s x <= 0 test. Both partition masks are built once (input.gt(0.0) and its logical_not) and reused, so their probabilities are exactly complementary by construction rather than by two independent comparisons.

ReLU

ReLU6

Parameter-free ReLU6: Hardtanh with fixed bounds 0 and 6. See Tensor::relu6 for the exact forward and (strictly-interior) backward semantics.

RecurrentConfig

Shared depth/direction/bias configuration for Rnn, Gru and Lstm. PyTorch’s RNNBase plays this role for its three subclasses, but it also owns shared weight storage and flattening behavior that Axis’s three families do not share (each keeps its own gate layout and init convention, matching LstmCell’s existing precedent of not introducing a second parameter convention). RecurrentConfig is the honest analogue: the three fields every family’s constructor actually shares, reused verbatim by Rnn::with_config, Gru::with_config and Lstm::with_config, with no shared class behavior invented beyond that.

ReflectionPad

Reflection padding, covering ReflectionPad1d/2d/3d. Matches PyTorch’s own constraint that each side’s padding be strictly less than the axis’s extent (checked before launch, in reflect_source’s doc); see edge_pad for the composition and its exact gradient.

RegimeSnapshot

ReplicationPad

Edge replication padding, covering ReplicationPad1d/2d/3d. No pad-size constraint beyond the shared overflow check: see edge_pad for the composition and its exact gradient.

ResidualLimits

ResidualScope

Declares what residual was evaluated, where, and by which numerical method.

RmsNorm

Root-mean-square normalization over one or more declared named axes.

Rnn

A stacked, optionally bidirectional, eager Elman RNN over one named time axis. See Lstm’s doc comment for the eager-transition contract and RecurrentConfig for num_layers/bidirectional/bias. The default configuration (one layer, unidirectional, bias on) is bit-exact with the single-layer RNN this type used to be.

RnnCell

One Elman RNN transition, h' = nonlinearity(W_ih x + b + W_hh h), with explicit hidden state. nonlinearity defaults to Tanh (PyTorch’s RNNCell(nonlinearity="tanh") default); .nonlinearity(Relu) matches nonlinearity="relu". Mirrors LstmCell’s single combined bias (rather than PyTorch’s separate bias_ih/bias_hh) and its Xavier-uniform initialization (rather than PyTorch’s uniform(-1/sqrt(hidden), 1/sqrt(hidden))); Axis’s own Lstm does not match PyTorch’s parameter layout or init either, so RnnCell follows LstmCell’s convention instead of introducing a second one. bias defaults to true (PyTorch’s default); .bias(false) matches bias=False.

RnnRun

Sequence output and the per-layer, per-direction terminal hidden states from one Rnn run, ordered as LstmRun::state documents.

SELU

Parameter-free SELU: PyTorch’s fixed-constant self-normalizing activation. See Tensor::selu for the exact constants.

SGD

Sample

Sequential

Shape

SiLU

Parameter-free sigmoid linear unit, x * sigmoid(x).

SignStraightThrough

Parameter-free sign quantization to {-1, +1} with a straight-through backward. See Tensor::sign_straight_through for the exact forward and backward semantics this reproduces from bae’s bitae.quantize.

SinglePass

A count-only invariant for a loader that promises unique examples within one pass. This proves the consumption budget, not the loader’s uniqueness promise.

Softmax2d

Softmax over an explicit named channel axis of a [channel, height, width] input (PyTorch’s Softmax2d, which requires exactly that rank). Delegates to the existing named-axis Tensor::softmax; height/width exist to state and check the contract, since softmax already treats every other axis independently.

Softmin

Softmin over an explicit named axis. See Tensor::softmin.

Softplus

Elementwise Softplus with explicit beta/threshold, defaulting to PyTorch’s own beta = 1, threshold = 20. See Tensor::softplus for the exact branch seam.

Softshrink

Parameter-free Softshrink with an explicit lambd. See Tensor::softshrink for the exact forward and gradient-at-the-band semantics.

Softsign

Parameter-free Softsign. See Tensor::softsign for the exact forward and backward semantics (smooth everywhere, no true kink).

Standardizer

A train-fitted affine standardization transform.

State

A persistent non-parameter tensor a module owns across training steps – BatchNorm-style running statistics, reported by Module::named_states. Cloning explicitly shares the same storage, like Parameter, but a State carries no gradient and is never touched by an optimizer: the only writer is TrainingPass::commit, which applies every update a forward_training call queued with TrainingPass::stage. Nothing ever writes a State outside that commit, so a forward pass alone – training or not – is side-effect free, and a failed training step (a panicking or error-returning loss, a failed backward or optimizer step) leaves every State exactly as it was.

Tanh

Parameter-free elementwise hyperbolic tangent module.

Tanhshrink

Parameter-free Tanhshrink. See Tensor::tanhshrink for the exact forward and backward semantics (smooth everywhere, no true kink).

Tensor

Immutable values. Clones share a graph node; detach starts a new, untracked leaf.

Threshold

Parameter-free Threshold with an explicit threshold/value. PyTorch requires both explicitly (no defaults); see Tensor::threshold for the exact forward and gradient-at-the-kink semantics.

TrainStep

Trainer

TrainingPass

The explicit per-step training context threaded through Module::forward_training. Trainer::step_training builds one from (run_seed, completed_steps); a caller driving forward_training outside the Trainer – a test, or a custom loop – builds one directly with TrainingPass::new.

Transformer

torch.nn.Transformer: one TransformerEncoder built from num_encoder_layers copies of a TransformerEncoderLayer plus a final LayerNorm, and one TransformerDecoder built the same way from num_decoder_layers copies of a TransformerDecoderLayer, sharing feature/time/embed_dim/num_heads/dim_feedforward and every dropout/activation/epsilon/norm_first/bias setting. forward encodes src into memory then decodes tgt against it, PyTorch’s own two-call composition. There is no custom_encoder/custom_decoder override: Axis has no Module-shaped container that fits a differently-typed encoder/decoder pair interchangeably here, so a fully custom pair is built directly from TransformerEncoder/TransformerDecoder instead.

TransformerDecoder

torch.nn.TransformerDecoder: the decoder-layer analogue of TransformerEncoder, sharing its independent-clone stacking convention and optional final LayerNorm.

TransformerDecoderLayer

One post-/pre-norm Transformer decoder block, torch.nn.TransformerDecoderLayer at the same PyTorch 2.14 defaults as TransformerEncoderLayer. Masked self-attention over tgt, then cross-attention (multihead_attn) whose query is tgt and whose key/value are memory, then the same feedforward block, with PyTorch’s own _sa_block/_mha_block/_ff_block inlined:

TransformerEncoder

torch.nn.TransformerEncoder: num_layers independently built copies of one configured, unbuilt TransformerEncoderLayer (cloned before build, mirroring PyTorch’s own _get_clones), plus an optional final LayerNorm. Each clone gets its own freshly seeded parameters at build time rather than PyTorch’s literal deep-copy-of-already-initialized weights (PyTorch’s _get_clones runs copy.deepcopy on one already-built layer, so every layer starts from bit-identical initial weights) – deliberately following Axis’s own stacking convention instead (every index-seeded container in this crate, e.g. Sequential::build, gives each position a distinct seed) and the design doc’s “Repeat must create independent blocks unless weight tying is explicit”.

TransformerEncoderLayer

One post-/pre-norm Transformer encoder block, torch.nn.TransformerEncoderLayer at PyTorch 2.14’s exact defaults (dim_feedforward=2048, dropout=0.1, activation="relu", layer_norm_eps=1e-5, norm_first=False, bias=True) and PyTorch’s own exact residual/dropout structure – _sa_block/_ff_block inlined here as private helpers of the same name:

Unflatten

Split one named axis into an ordered list of named axes whose extents multiply back to it, PyTorch’s nn.Unflatten over identities instead of a positional dim and an unnamed size tuple. A thin Module wrapper over Tensor::split; see its own doc for the exact validation and physical layout it reuses verbatim.

Unfold

Named-axis im2col (PyTorch’s nn.Unfold): extracts sliding local kernel-sized blocks, replacing channels in place with patch (extent channels * kernel volume) and shrinking spatial to their windowed count. Out-of-bounds (padding) positions read 0.0. Stride defaults to [1, 1] and padding to [0, 0]. Stateless: build only validates and infers the output shape.

Upsample

torch.nn.Upsample(size=None, scale_factor=None, mode='nearest', align_corners=None) over a fixed, ordered list of named spatial axes (PyTorch infers 1D/2D/3D from the input’s trailing dims; here the axes are named explicitly instead, so any rank works). Build with Upsample::size (a literal per-axis output extent, PyTorch’s size=) or Upsample::scale_factor (a per-axis multiplier, PyTorch’s scale_factor=) – exactly one form, matching PyTorch’s own mutual exclusivity.

ZeroPad

Exact zero padding, covering ZeroPad1d/ZeroPad2d/ZeroPad3d (see the tuple-to-axis mapping above validate_pad_list). Composed entirely from repeated Tensor::pad_zeros, one call per listed axis; backward is pad_zeros’s own exact crop, so a source element’s gradient is unaffected by padding.

Enums

AttentionMask

An additive (f32, added to pre-softmax scores; -infinity or a large negative forbids a query/key pair) or boolean attn_mask for MultiheadAttention. The boolean variant uses PyTorch’s own convention: a 1.0 entry forbids that query/key pair from attending, 0.0 allows it – the opposite sense from Tensor::masked_softmax’s own validity mask, which MultiheadAttention inverts internally before calling it.

DataRegimeReceipt

DisjointnessEvidence

The strength of the evidence recorded by a disjointness receipt.

EmbeddingBagMode

Pooling mode for EmbeddingBag, matching PyTorch’s mode argument.

KernelSpec

Physical kernel shapes for crate::Device::prepare_kernels.

LearningDirection

LearningEvidence

The strength of the claim made by a learning-progress receipt.

MonotoneDirection

MuonMatrixOrientation

The storage orientation of one explicitly selected Muon matrix.

PopulationMode

Whether the ledger keeps identities after checking an observation.

ResidualClaim

The claim made by this receipt. A guaranteed law requires a separate, construction-backed API and cannot be selected by a caller here.

ResidualEvidence

RnnNonlinearity

The pointwise nonlinearity applied at each Rnn/RnnCell transition. PyTorch’s default is Tanh; Relu matches nonlinearity="relu".

TransformerActivation

Encoder/decoder feedforward nonlinearity, torch.nn.TransformerEncoderLayer’s activation argument restricted to its two string spellings ("relu"/"gelu"; a custom callable has no analogue here). PyTorch’s own _get_activation_fn maps "gelu" to F.gelu at its default approximate="none" – the exact erf form, Tensor::gelu_exact / ExactGELU in this crate, not the tanh-form GELU module.

UpsampleMode

torch.nn.Upsample’s interpolation mode (mode=).

Traits

DataSource

A source may be finite or generated. None availability means no finite corpus exists.

IntoAxes

IntoLayers

Module

Optimizer

Functions

clip_grad_norm

PyTorch’s clip_grad_norm_ (https://docs.pytorch.org/docs/2.14/generated/torch.nn.utils.clip_grad_norm_.html): one global L2 norm over every listed parameter’s accumulated gradient together, not a per-tensor norm. Every gradient is rescaled by the SAME factor max_norm / (total_norm + 1e-6), and only when total_norm exceeds max_norm; below the threshold every gradient is left bit-exact, never renormalized to exactly max_norm. Returns the pre-clip total norm either way, matching PyTorch’s own return value.

cosine_annealing_lr

Cosine-annealed learning rate at step (0-indexed), matching PyTorch’s CosineAnnealingLR closed form (https://docs.pytorch.org/docs/2.14/generated/torch.optim.lr_scheduler.CosineAnnealingLR.html):

jit_cache_stats

Snapshot of the cuTile JIT disk cache’s cumulative hit/miss counts for this process. Always zero in the docs.rs build, which has no CUDA runtime and never creates a device.

one_cycle_lr

One-cycle learning rate at step (0-indexed of total_steps), matching PyTorch’s OneCycleLR at its own defaults (https://docs.pytorch.org/docs/2.14/generated/torch.optim.lr_scheduler.OneCycleLR.html): a cosine warmup from max_lr / div_factor up to max_lr over the first pct_start fraction of total_steps, then a cosine anneal down to max_lr / div_factor / final_div_factor. PyTorch’s own closed form (torch/optim/lr_scheduler.py, _annealing_cos and OneCycleLR.get_lr):

Type Aliases

Result