Struct Tensor 

pub struct Tensor(/* private fields */);

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

Implementations

impl Tensor

pub fn categorical_accuracy( &self, targets: &Tensor, class: Axis, ) -> Result<CategoricalAccuracy>

Compare categorical argmaxes on the device and return exact row counts. targets may contain one-hot or probability distributions.

pub fn masked_categorical_accuracy( &self, targets: &Tensor, mask: &Tensor, class: Axis, ) -> Result<CategoricalAccuracy>

Compare categorical argmaxes only where mask is one.

The mask must have the logits’ shape with the class axis removed and contain only zeroes and ones. Returning counts keeps evaluation exactly composable across batches.

impl Tensor

pub fn from_slice( values: &[f32], dims: impl IntoIterator<Item = Dim>, device: &Device, ) -> Result<Self>

pub fn zeros( dims: impl IntoIterator<Item = Dim>, device: &Device, ) -> Result<Self>

Allocate a device-resident zero tensor without constructing a host vector.

pub fn uniform( dims: impl IntoIterator<Item = Dim>, seed: u64, low: f32, high: f32, device: &Device, ) -> Result<Self>

Deterministic uniform draw in [low, high) from the shared xorshift stream that initializes parameters, generated host-side then uploaded like Tensor::from_slice. A random draw has no upstream input, so the result carries no gradient edge; it is a constant, not a parameter. The same seed, shape and range reproduce identical values on any run or machine; distinct seeds diverge. Inherits the shared stream’s documented quirk: a seed below 2^40 draws exactly low first.

pub fn normal( dims: impl IntoIterator<Item = Dim>, seed: u64, mean: f32, std: f32, device: &Device, ) -> Result<Self>

Deterministic normal draw with the given mean and standard deviation std, from the same shared xorshift stream as Tensor::uniform via the Box-Muller transform: two raw [0, 1) stream samples u1, u2 become one pair z0 = sqrt(-2 * ln(1 - u1)) * cos(2*pi*u2), z1 = sqrt(-2 * ln(1 - u1)) * sin(2*pi*u2) of independent standard-normal values, each scaled to mean + std * z; an odd element count drops the unused second value of the final pair. Generated host-side then uploaded, with no gradient edge, exactly like Tensor::uniform.

pub fn shape(&self) -> &Shape

pub fn extent(&self, axis: Axis) -> Result<usize>

pub fn device(&self) -> &Device

pub fn requires_grad(&self) -> bool

pub fn with_grad(&self) -> Self

pub fn detach(&self) -> Self

pub fn to_vec(&self) -> Result<Vec<f32>>

Values in logical Shape order, independent of physical storage order.

pub fn item(&self) -> Result<f32>

pub fn grad(&self) -> Option<Self>

Leaf gradients accumulate until explicitly cleared. Intermediates do not retain gradients.

pub fn zero_grad(&self)

pub fn add(&self, rhs: &Self) -> Result<Self>

pub fn sub(&self, rhs: &Self) -> Result<Self>

pub fn mul(&self, rhs: &Self) -> Result<Self>

pub fn div(&self, rhs: &Self) -> Result<Self>

Elementwise self / rhs. Same shape-agreement contract as Tensor::mul: identical axis sets, no implicit alignment. Division by zero follows IEEE float semantics (produces inf/nan, as in PyTorch); callers that need a safe denominator must clamp it themselves before dividing.

pub fn squared_error(&self, rhs: &Self) -> Result<Self>

pub fn abs(&self) -> Result<Self>

Elementwise absolute value. Backward is gradient * sign(x), matching PyTorch’s abs backward: the gradient is exactly zero at x == 0 (PyTorch’s own subgradient choice there, not NaN or either one-sided slope).

pub fn absolute_error(&self, rhs: &Self) -> Result<Self>

Elementwise mean absolute error (PyTorch’s unreduced F.l1_loss): (self - rhs).abs(). Same axis-agreement contract as Tensor::squared_error (identical axis sets, equal extents, unreduced shape), leaving the reduction (typically .mean(axes)) to the caller.

pub fn huber_loss(&self, rhs: &Self, delta: f32) -> Result<Self>

PyTorch-exact HuberLoss (default delta = 1.0): 0.5 * (self - rhs)^2 where |self - rhs| < delta, else delta * (|self - rhs| - 0.5 * delta). Composed as 0.5 * min(|d|, delta)^2 + delta * relu(|d| - delta) with d = self - rhs – an exact algebraic identity with PyTorch’s two-branch definition, continuous and differentiable at |d| == delta – so no dedicated kernel or Rule is needed; backward is clamp(d, -delta, delta), the composition’s own derivative. Same axis-agreement contract as Tensor::squared_error/Tensor::absolute_error: identical axis sets, unreduced shape, reduction left to the caller. delta must be finite and positive.

pub fn smooth_l1_loss(&self, rhs: &Self, beta: f32) -> Result<Self>

PyTorch-exact SmoothL1Loss (default beta = 1.0), the family’s own scaling of Tensor::huber_loss: SmoothL1Loss(beta) == HuberLoss(delta = beta) / beta is an exact algebraic identity (both branches divide out), so this composes huber_loss rather than repeating its formula. beta must be finite and positive; morpheus’s RBC/WBC radius and offset regression heads call F.smooth_l1_loss(..., beta=.02) throughout research/src/vision/morpheus/mobilesam/scale10 (for example train_click_rbc.py:558), the consumer for the non-default beta this test exercises.

pub fn nll_loss(&self, targets: &Self, class: Axis) -> Result<Self>

NLLLoss generalized to constant one-hot or probability targets over a named class axis, exactly the way Tensor::categorical_cross_entropy_with_logits generalizes CrossEntropyLoss – the same target representation decision, mirrored here. self must already hold log-probabilities (typically logits.softmax(class)?.ln()? or a named-axis log-softmax composition); unlike categorical_cross_entropy_with_logits, this op folds in no softmax of its own. The named class axis is reduced; every other axis is preserved, and a target may omit axes self has (broadcasting one target row over them, as categorical_cross_entropy_with_logits also allows). Loss is -sum(class, target * self), so backward is exactly -target, the composition’s own linear derivative – no dedicated kernel or Rule. morpheus’s partition classifier calls F.nll_loss(q[full].log(), lab[full].long(), weight=ce_weight) (research/src/vision/morpheus/mobilesam/scale10/train_partition5.py:291); Axis has no per-class weight yet. Targets are constants in reverse mode.

pub fn nll_loss_weighted( &self, targets: &Self, class: Axis, weight: &Self, ignore_index: Option<usize>, ) -> Result<Self>

Tensor::nll_loss with PyTorch’s two remaining NLLLoss options: a per-class weight over the named class axis, and an ignore_index. weight follows the per-class-weight convention Tensor::binary_cross_entropy_with_logits_weighted already uses for pos_weight: it must contain only the class axis (or a subset that still broadcasts against it), one value per class. Weighting multiplies each class’s term before the class axis is summed, -sum(class, weight * target * self), which is exactly PyTorch’s weight[c] * log_prob[c] at a one-hot target and generalizes linearly to a soft one. Unweighted nll_loss is the special case weight == 1.

PyTorch’s reduction='mean' (which this op, like every other unreduced loss here, leaves to the caller’s own mean/sum) does NOT divide by the row count: it divides by the SUM OF THE SELECTED WEIGHTS – weight[target[n]] for every kept row – not by N. A caller reproducing PyTorch’s mean must compute that weight sum itself (for a one-hot target, weight[target] is targets.mul(weight)?.sum(class)) and divide the summed unreduced loss by it, rather than calling the ordinary axis mean.

ignore_index, given as Some(class_index), zeroes a row’s contribution in proportion to how much target probability mass sits on that class: loss_row *= 1 - target[ignore_index]. At a one-hot target this is PyTorch’s exact all-or-nothing ignore; a soft target scales down partially rather than erroring, the same generalization Tensor::nll_loss itself already makes for one-hot targets. An ignored row must also be excluded from the weight sum above when reproducing PyTorch’s mean, for the same reason. ignore_index must be a valid class index.

pub fn binary_cross_entropy(&self, targets: &Self) -> Result<Self>

Elementwise BCELoss for probability inputs: -[y * log(x) + (1 - y) * log(1 - x)]. self holds probabilities in [0, 1] (already squashed by the caller, typically .sigmoid()) – unlike Tensor::binary_cross_entropy_with_logits, this op fuses no sigmoid. Matching PyTorch exactly, each logarithm is floored at PyTorch’s own literal -100 (f32::max(log(x), -100.0) / f32::max(log(1 - x), -100.0)), so a prediction of exactly 0 or 1 gives a finite loss rather than NaN. Earlier this op composed the same floor from Tensor::clamp/Tensor::ln instead, applied to the input at f32::MIN_POSITIVE (the closest a differentiated-through-ln composition could get to -100, since flooring ln’s own output would still let its backward see g / 0 -> NaN at x == 0, and flooring the input at PyTorch’s exp(-100) – an f32 subnormal – makes ln’s backward 1 / x overflow to infinity, which the clamp’s own zeroing multiplier then turns into inf * 0 -> NaN right back). This dedicated kernel pair (bce_loss/bce_loss_backward) breaks that dependency: forward floors the output of each logarithm at exactly PyTorch’s -100, and backward is PyTorch’s own explicit formula, (x - y) / max((1 - x) * x, eps) with eps = 1e-12, computed directly rather than falling out of differentiating through the floored ln. It agrees with PyTorch’s interior formula (x - y) / (x * (1 - x)) away from the eps floor, and stays finite (rather than NaN) at x == 0 and x == 1. Same axis-agreement contract as Tensor::absolute_error. Targets are constants in reverse mode.

pub fn kl_div_loss(&self, targets: &Self) -> Result<Self>

Elementwise KLDivLoss at PyTorch’s default log_target = false: self holds log-probabilities (as NLLLoss/CrossEntropyLoss do), target holds probabilities, and the pointwise loss is target * (log(target) - self). Matches PyTorch’s xlogy convention at target == 0 – contributes exactly 0, never NaN from 0 * -inf – by substituting 1.0 for log’s input only where target == 0, a substitution whose own target factor of 0 cancels it either way. log_target = true is not implemented. Composed from Tensor::eq/Tensor::logical_not/Tensor::ln; no dedicated kernel. morpheus’s distillation heads call F.kl_div(F.log_softmax(logits / t, -1), F.softmax(teacher_logits / t, -1), ...) (research/src/vision/morpheus/mobilesam/scale10/microtier/runs/wbc-edgepath-slice4m/train.py:124), the consumer for this exact log_target = false shape. Targets are constants in reverse mode.

pub fn kl_div_loss_log_target(&self, targets: &Self) -> Result<Self>

Tensor::kl_div_loss at PyTorch’s log_target = true: targets also holds log-probabilities (rather than probabilities), and the pointwise loss is exp(target) * (target - self) – PyTorch’s own log_target = true formula. Unlike log_target = false, no xlogy zero-target substitution is needed: target never appears inside a logarithm here (only self does, and self is the differentiable operand, not something this op takes the log of), so there is no 0 * -inf case to guard. Composed from Tensor::exp/Tensor::sub/Tensor::mul; no dedicated kernel. Targets are constants in reverse mode.

pub fn poisson_nll_loss(&self, targets: &Self) -> Result<Self>

PoissonNLLLoss at PyTorch’s default log_input = true, full = false: self holds the log of the predicted rate, and the pointwise loss is exp(self) - target * self. Composed from Tensor::exp; no dedicated kernel. Targets are constants in reverse mode. Delegates to Tensor::poisson_nll_loss_full with full = false, so this remains bit-exact with the earlier, narrower implementation.

pub fn poisson_nll_loss_full( &self, targets: &Self, log_input: bool, full: bool, eps: f32, ) -> Result<Self>

Tensor::poisson_nll_loss with PyTorch’s remaining PoissonNLLLoss options.

log_input: when true (PyTorch’s default), self holds the log of the predicted rate and the base loss is exp(self) - target * self, exactly Tensor::poisson_nll_loss. When false, self holds the rate itself and the base loss is self - target * log(self + eps), PyTorch’s own log_input = false formula; eps (must be finite and positive) keeps the logarithm finite at self == 0, matching PyTorch’s own stabilization. eps is unused when log_input = true, matching PyTorch.

full: when true, adds PyTorch’s Stirling approximation term, target*log(target) - target + 0.5*log(2*pi*target), but only where target > 1 (PyTorch’s own threshold; the term is 0 at target <= 1, including target == 0, where the term’s own logarithm would otherwise be undefined). The masked branch is computed the same xlogy-style substitution Tensor::kl_div_loss already uses: a safe stand-in target of 1.0 (whose own Stirling term is exactly 0) is substituted wherever the mask is 0, so no element’s logarithm ever sees a non-positive input before the mask zeroes its contribution. The Stirling term does not depend on self, so it contributes no gradient; composed entirely from Tensor::gt/Tensor::ln/ Tensor::mul/Tensor::add; no dedicated kernel.

pub fn gaussian_nll_loss( &self, targets: &Self, var: &Self, eps: f32, ) -> Result<Self>

GaussianNLLLoss at PyTorch’s default full = false: self is the predicted mean, var the predicted variance (itself differentiable – unlike targets, matching PyTorch treating var as a second model output, not a label), and the pointwise loss is 0.5 * (log(max(var, eps)) + (self - target)^2 / max(var, eps)), composed from Tensor::clamp/Tensor::ln/Tensor::squared_error/Tensor::div with no dedicated kernel. var must be elementwise nonnegative before clamping, matching PyTorch’s own check; eps must be finite and positive. Unlike PyTorch, which clamps var inside no_grad so its gradient passes straight through the clamp using the clamped value, Axis composes this from the ordinary clamp, whose ordinary boundary rule zeroes var’s own gradient at elements the clamp actually moved (var < eps); self’s gradient is unaffected either way. Targets are constants in reverse mode. See Tensor::gaussian_nll_loss_full for PyTorch’s full = true constant term.

pub fn gaussian_nll_loss_full( &self, targets: &Self, var: &Self, eps: f32, full: bool, ) -> Result<Self>

Tensor::gaussian_nll_loss with PyTorch’s remaining full option: adds the constant 0.5 * log(2*pi) term PyTorch’s full = true includes (the normalizing term of a Gaussian log-likelihood that full = false omits as an additive constant with no gradient). Depends on neither self, targets, nor var, so it contributes no gradient to any of them; composed as Tensor::gaussian_nll_loss plus that one scalar constant, reusing every check and the exact full = false value rather than recomputing it. full = false is bit-exact with Tensor::gaussian_nll_loss itself.

pub fn soft_margin_loss(&self, targets: &Self) -> Result<Self>

Elementwise SoftMarginLoss: log(1 + exp(-target * self)) for target in {-1, +1}, composed as (-target * self).softplus(1.0, 20.0) – PyTorch’s own default Softplus threshold, reused here purely for the numerically stable linear seam, not as a knob callers choose. Backward is the composition’s own -target * sigmoid(-target * self); no dedicated kernel. Same axis-agreement contract as Tensor::absolute_error. Targets are constants in reverse mode.

pub fn multilabel_soft_margin_loss( &self, targets: &Self, class: Axis, ) -> Result<Self>

MultiLabelSoftMarginLoss at PyTorch’s default (no per-class weight): the mean, over the named class axis, of Tensor::binary_cross_entropy_with_logits. PyTorch’s own formula -1/C * sum_c [y_c*log(sigmoid(x_c)) + (1-y_c)*log(1-sigmoid(x_c))] is exactly that mean, since binary_cross_entropy_with_logits already computes the summand elementwise – no dedicated kernel or Rule is needed. Unlike Axis’s other losses, PyTorch’s own unreduced (reduction='none') form already reduces the class axis, so the class axis is gone from the result before any caller reduction; every other axis is preserved. Inherits binary_cross_entropy_with_logits’s constant-target requirement and identical-axis-set contract.

pub fn binary_cross_entropy_with_logits(&self, targets: &Self) -> Result<Self>

Stable elementwise binary cross-entropy. Targets are constants in reverse mode.

pub fn binary_cross_entropy_with_logits_weighted( &self, targets: &Self, pos_weight: &Self, ) -> Result<Self>

Stable elementwise binary cross-entropy with an explicit positive-class weight, matching torch.nn.BCEWithLogitsLoss(pos_weight=...): the positive term of each element’s loss is scaled by pos_weight before the negative term is added, so pos_weight changes the loss value itself and is not reachable by scaling the unweighted loss afterward. pos_weight may be a scalar or may broadcast over any subset of self’s axes (for example one weight per class, broadcasting over batch and spatial axes); it follows the same subset-axis broadcasting as elementwise add/mul, never an implicit outer product. Targets and pos_weight are constants in reverse mode.

pub fn categorical_cross_entropy_with_logits( &self, targets: &Self, class: Axis, ) -> Result<Self>

Stable categorical cross-entropy for constant one-hot or probability targets. The named class axis is reduced; every other logits axis is preserved.

impl Tensor

pub fn linear_cross_entropy_with_logits( &self, weight: &Self, bias: Option<&Self>, input_axis: Axis, targets: &Self, class: Axis, options: LinearCrossEntropyOptions, ) -> Result<Self>

torch.nn.functional.linear_cross_entropy: composes the final linear projection – self.contract(weight, input_axis) plus an optional bias, exactly the body of [crate::model::nn::Linear]’s own forward – with Tensor::categorical_cross_entropy_with_logits. weight is [input_axis, class] (PyTorch’s linear_weight is [C, in_features]; the axis carrying class here is read off weight itself, so it must equal class), bias is [class].

PyTorch’s real linear_cross_entropy is a FUSED kernel whose documented purpose is to avoid ever materializing the [*, C] logits tensor, roughly halving peak activation memory at a large vocabulary C (its stated motivation: LLM heads with C in the 128-256K range). This composition produces bit-identical forward values and an exact gradient through ordinary reverse-mode autodiff, but it materializes the full logits tensor like any ordinary Linear + CrossEntropyLoss pair – it does NOT capture the fused kernel’s memory-saving contract. Graded PARTIAL for exactly this reason; see the landing PR for the full accounting.

label_smoothing (PyTorch default 0.0, must be in [0, 1)) blends targets toward the uniform distribution over class’s width classes before the cross-entropy call: (1 - a) * target + a / width, PyTorch’s own documented formula.

pub fn adaptive_log_softmax_with_loss( &self, input_axis: Axis, targets: &[usize], head_weight: &Self, head_bias: Option<&Self>, tail_weights: &[(Self, Self)], cutoffs: &[usize], n_classes: usize, div_value: f32, ) -> Result<Self>

torch.nn.AdaptiveLogSoftmaxWithLoss(in_features, n_classes, cutoffs, div_value=4.0, head_bias=False)’s forward pass (Grave et al., “Efficient softmax approximation for GPUs”), composed from self.contract (the head and per-cluster tail projections, each exactly Linear’s own body), Tensor::log_softmax, and one-hot selection built as a host constant then reduced with mul + sum – the same masked-selection idiom categorical_cross_entropy_with_logits itself already uses, substituted for Tensor::gather because gather applies one host index list uniformly across every row, while each row here targets a different class.

input is [batch_axis, input_axis] and targets is one class index (< n_classes) per row of batch_axis, matching PyTorch’s own (N, in_features) / (N,) restriction (AdaptiveLogSoftmaxWithLoss does not generalize to extra batch axes either). cutoffs are the tail cluster boundaries (PyTorch’s own constructor argument, NOT including n_classes): the shortlist is class range [0, cutoffs[0]), projected straight through head_weight’s first cutoffs[0] columns; tail cluster c (0..cutoffs.len()) is class range [edges[c], edges[c + 1]) where edges = [cutoffs[0], cutoffs[1], .., n_classes], reached through head_weight’s cutoffs[0] + c cluster-indicator column plus its own two-stage tail_weights[c] projection ((down, up), down: [input_axis, hidden], up: [hidden, tail_class], no bias, exactly PyTorch’s nn.Sequential tail). div_value (default 4.0) is validated against each down’s hidden extent, PyTorch’s own formula hidden = floor(in_features / div_value ** (c + 1)), rejecting a mis-shaped tail projection before launch. head_bias is Some for PyTorch’s head_bias=True, None for the documented default head_bias=False.

Returns the scalar reduction='mean' negative log-likelihood loss (PyTorch’s own forward return is (output, loss); output, the per-row target log-probability, and the separate log_prob/predict full-distribution queries are NOT implemented here – an honest gap, not attempted, which is why this row is graded PARTIAL rather than YES).

pub fn ctc_loss( log_probs: &Self, time_axis: Axis, class_axis: Axis, targets: &[Vec<usize>], blank: usize, ) -> Result<Self>

torch.nn.CTCLoss(blank=0, reduction='mean', zero_infinity=False)’s forward pass, over the log-space forward recursion (Graves 2006): build the length-2L + 1 blank-interleaved extended target sequence per row, then for each time step t (a host loop, driven from the host per time step as directed) combine the previous step’s alpha[t-1] with itself shifted by one and by two along the extended axis via Tensor::logsumexp, add that step’s emission log-probability, and repeat. Every step is an ordinary composition of concat/narrow (the shift), stack + logsumexp (the pairwise/triple log-sum), mul + sum against a host-built one-hot selection (the emission lookup – same substitution for gather as Tensor::adaptive_log_softmax_with_loss, since each row’s target sequence differs), and add. No hand-derived backward is coded: the forward recursion alone is a complete, differentiable definition of -log P(target | input), so ordinary reverse-mode autodiff through this graph supplies CTC’s textbook beta-based gradient exactly, for free, the same way resample_bilinear’s contract call gets an exact backward without a dedicated Rule.

log_probs is [time_axis, batch_axis, class_axis], already log-probabilities (Tensor::log_softmax(class_axis), matching PyTorch’s own documented input contract). targets is one target sequence per row of batch_axis, ALL of the same length L (targets[row].len()), and time_axis’s full extent T is treated as every row’s input length. This is an honest, named restriction, not PyTorch’s general contract: real CTCLoss takes separate input_lengths/target_lengths per row so a batch can mix padded sequences of different lengths; that per-row length masking is NOT implemented here. blank is the blank class index (PyTorch default 0); zero_infinity (clamping an infinite loss/gradient to 0, PyTorch default False) is not implemented, matching the default. reduction='mean': PyTorch’s own documented mean divides each row’s loss by its target length THEN means over the batch; with L shared here that is mean(loss_row) / L.

pub fn scale(&self, factor: f32) -> Result<Self>

pub fn relu(&self) -> Result<Self>

pub fn sigmoid(&self) -> Result<Self>

Stable elementwise logistic sigmoid.

pub fn silu(&self) -> Result<Self>

Elementwise sigmoid linear unit, x * sigmoid(x).

pub fn leaky_relu(&self, negative_slope: f32) -> Result<Self>

Elementwise leaky ReLU with an explicit finite, non-negative slope. Its derivative at exactly zero is the negative slope, matching PyTorch.

pub fn tanh(&self) -> Result<Self>

Elementwise hyperbolic tangent.

pub fn sin(&self) -> Result<Self>

Elementwise sine in radians.

pub fn gelu(&self) -> Result<Self>

Tanh-approximated GELU, matching the common transformer formulation.

This is PyTorch’s F.gelu(x, approximate="tanh"), NOT its default. A model ported from an un-annotated F.gelu(x) or nn.GELU() wants Self::gelu_exact; using this form instead gives plausible numbers that differ by about 1e-3 in the tails.

pub fn gelu_exact(&self) -> Result<Self>

Erf-form GELU, x * Phi(x), rather than the tanh formulation of Self::gelu.

The CUDA backend numerically evaluates the normal CDF in FP32 (not bitwise identical to a particular libm). Backward uses Phi(x) + x * phi(x). This distinction matters when importing pretrained exact-GELU models: this is PyTorch’s default F.gelu/nn.GELU (approximate="none").

pub fn sign_straight_through(&self) -> Result<Self>

Elementwise sign quantization to {-1, +1} with a straight-through backward.

Forward: +1 where x > 0, otherwise -1, so x == 0 maps to -1. This reproduces bae’s bitae.quantize two-level branch, torch.where(z > 0, 1.0, -1.0), rather than torch.sign, which would give exactly 0 at x == 0. Backward passes the upstream gradient through unchanged (the straight-through estimator, Rule::Identity), matching z + (hard - z).detach() in the same reference: the hard threshold has no gradient of its own, so the whole local Jacobian is the identity. Output has the same Shape as the input.

pub fn exp(&self) -> Result<Self>

Elementwise exponential, exp(x). Backward is g * exp(x), computed from the already-produced output rather than re-evaluating exp from the input, the same trade Self::tanh makes. Output has the same Shape as the input.

pub fn ln(&self) -> Result<Self>

Elementwise natural logarithm, ln(x). IEEE behaviour, no clamping: ln of a non-positive input is -inf at exactly x == 0 and NaN for x < 0, matching f32::ln/PyTorch’s torch.log rather than any epsilon or absolute-value guard. Backward is g / x, which inherits the same non-finite behaviour at and below zero. Output has the same Shape as the input.

pub fn softplus(&self, beta: f32, threshold: f32) -> Result<Self>

PyTorch-exact Softplus: (1 / beta) * ln(1 + exp(beta * x)), except where beta * x > threshold, which returns x itself (the identity function’s large-x asymptote, computed exactly rather than through the logarithm) to keep both branches finite and match torch.nn.functional.softplus(x, beta, threshold) bit for bit at the seam. Backward is sigmoid(beta * x) on the logarithmic branch and exactly 1 on the linear branch, mirroring the forward’s own branch selection. beta must be finite and positive; threshold must be finite.

pub fn elu(&self, alpha: f32) -> Result<Self>

Elementwise ELU (torch.nn.functional.elu): x where x > 0, otherwise alpha * (exp(x) - 1). alpha must be finite and nonzero – PyTorch’s own documented domain places no sign restriction on alpha (torch.nn.ELU’s docs give only “the alpha value for the ELU formulation. Default: 1.0”, and the piecewise formula itself is applied as written for any nonzero alpha, negative included: unlike celu below there is no max(0, x) + min(0, ...) wrapper in PyTorch’s ELU definition, so the negative branch is used exactly as alpha * (exp(x) - 1) regardless of alpha’s sign). The positive branch (x itself) never reads alpha, so every already-recorded alpha > 0 output and gradient stays bit-exact; only the guard widens. Composed from Self::gt/Self::logical_not, Self::clamp and Self::exp rather than a dedicated kernel: clamping the input to (-inf, 0] before exp keeps the positive branch – masked out of the result anyway – from ever overflowing. Backward matches PyTorch’s own x > 0 split exactly, including at x == 0 (grouped with the negative branch, so the gradient there is alpha, not 1): clamp’s own pass-through gradient at its upper bound keeps that boundary case exact through the chain rule.

pub fn celu(&self, alpha: f32) -> Result<Self>

Elementwise CELU (torch.nn.functional.celu, the continuously differentiable exponential linear unit): PyTorch’s documented max(0, x) + min(0, alpha * (exp(x / alpha) - 1)), “valid for alpha != 0” per torch.nn.CELU’s own docs, negative alpha included. Composed as x where x > 0, otherwise alpha * (exp(x / alpha) - 1), exactly like Self::elu’s two-branch split, WITHOUT an explicit min(0, ...)/max(0, ...): for x > 0 and any nonzero alpha, x / alpha and alpha always land on opposite sides of the sign that makes alpha * (exp(x / alpha) - 1) >= 0 (so min(0, ...) collapses to 0 and the sum is exactly x), and for x <= 0 the same sign algebra makes alpha * (exp(x / alpha) - 1) land on the correct sign to equal min(0, ...) unclamped – so the two-branch form already equals PyTorch’s max/min form for every nonzero alpha, not only positive ones. The positive branch never reads alpha, so every already-recorded alpha > 0 output and gradient stays bit-exact; only the guard widens. Scales the clamped input by 1 / alpha before exp and the branch result back up by alpha; the same x == 0 boundary reasoning as elu applies.

pub fn selu(&self) -> Result<Self>

Elementwise SELU (torch.nn.functional.selu): PyTorch’s fixed-constant self-normalizing activation, scale * elu(x, alpha) with alpha = 1.6732632423543772 and scale = 1.0507009873554805 – PyTorch’s own literals, solved so a standard-normal input keeps zero mean and unit variance through the nonlinearity. No configurable parameters: unlike Self::elu/Self::celu, selu takes no alpha.

pub fn log_sigmoid(&self) -> Result<Self>

Numerically stable elementwise log-sigmoid, ln(sigmoid(x)) (torch.nn.functional.logsigmoid), computed as -softplus(-x): PyTorch’s own stabilization. softplus’s existing linear seam keeps this finite for very negative x, where a literal sigmoid then ln would underflow to ln(0) = -inf, and it asymptotes to x exactly as log_sigmoid should when x -> -inf. Backward is sigmoid(-x), which falls out of the chain rule through Self::scale and Self::softplus with no dedicated rule.

pub fn mish(&self) -> Result<Self>

Elementwise Mish (torch.nn.functional.mish), x * tanh(softplus(x)), composed from the existing stable Self::softplus (beta = 1, threshold = 20, PyTorch’s own defaults) and Self::tanh rather than a dedicated kernel.

pub fn glu(&self, axis: Axis) -> Result<Self>

Gated Linear Unit (torch.nn.functional.glu): split axis into two equal halves and gate the first half by the sigmoid of the second, a * sigmoid(b). axis’s extent must be even; both halves keep axis’s own identity, at half the extent. Composed from Self::narrow and Self::sigmoid/Self::mul.

pub fn prelu(&self, weight: &Self) -> Result<Self>

Elementwise parametric ReLU (torch.nn.functional.prelu): x where x > 0, otherwise weight * x. weight is a constant here – the crate::PReLU module supplies a differentiable crate::Parameter and composes with this method – and follows the same subset-axis broadcasting as Self::mul, so a single shared weight (shape []) or one weight per entry of a named channel axis both work unchanged. Composed entirely from Self::gt/Self::logical_not and Self::mul/Self::add, with the same x > 0 boundary convention as Self::elu; weight’s own gradient (when it requires one) falls out of mul’s existing broadcast-sum backward with no dedicated rule.

pub fn log_softmax(&self, axis: Axis) -> Result<Self>

Numerically stable log-softmax along one named axis (torch.nn.functional.log_softmax), x - logsumexp(x, axis), composed entirely from the existing Self::logsumexp (itself max-shifted) and Self::sub’s broadcast over the axis logsumexp removes – no dedicated kernel or backward rule. Shares logsumexp’s all-non-finite group convention (NaN, not PyTorch’s -infinity).

pub fn softmin(&self, axis: Axis) -> Result<Self>

Softmin along one named axis (torch.nn.functional.softmin), softmax(-x, axis), composed from Self::scale and the existing stable Self::softmax.

pub fn gt(&self, scalar: f32) -> Result<Self>

Elementwise self > scalar as a {0.0, 1.0} mask, IEEE-ordered so any comparison against NaN is false (0.0), exactly like PyTorch’s >. No autograd edge, even when self requires grad: like PyTorch, a comparison is not differentiable.

pub fn ge(&self, scalar: f32) -> Result<Self>

Elementwise self >= scalar. Same NaN and no-gradient contract as Self::gt.

pub fn lt(&self, scalar: f32) -> Result<Self>

Elementwise self < scalar. Same NaN and no-gradient contract as Self::gt.

pub fn le(&self, scalar: f32) -> Result<Self>

Elementwise self <= scalar. Same NaN and no-gradient contract as Self::gt.

pub fn eq(&self, scalar: f32) -> Result<Self>

Elementwise self == scalar, bit-for-bit IEEE equality (no epsilon). Same NaN and no-gradient contract as Self::gt.

pub fn logical_and(&self, rhs: &Self) -> Result<Self>

Elementwise logical AND of two {0.0, 1.0} masks, spelled as their product (mask * mask is exactly PyTorch’s & on 0/1 tensors). Same axis-agreement contract as Self::mul; the result carries a gradient only if mul would give one, which is never the case for two comparison-derived masks since neither operand has an autograd edge to begin with.

pub fn logical_not(&self) -> Result<Self>

Elementwise logical NOT of a {0.0, 1.0} mask, spelled as 1.0 - mask (PyTorch’s ~ on a 0/1 tensor). Logical OR has no named method because no migrated consumer calls it; on 0/1 masks it is an elementwise maximum (a | b == max(a, b), since both are {0.0, 1.0}), which Axis has no tensor/tensor primitive for yet – add one only when a consumer needs it.

pub fn inverse_sqrt(&self, epsilon: f32) -> Result<Self>

Elementwise (x + epsilon)^-1/2; inputs plus epsilon must be positive.

pub fn clamp(&self, min: Option<f32>, max: Option<f32>) -> Result<Self>

Elementwise clamp into [min, max]; either bound may be None to leave that side unbounded, matching torch.clamp’s optional min/max keywords (sites.clamp_(0, 1) in vision/image-encode’s Stage A and counts.clamp(min=1) in world/energy-output’s reconstruction loss are the two consumers – the second gives only min). Rejects a NaN bound or min > max before any device launch. A NaN element of self propagates unclamped: like Self::gt and its siblings, every ordered comparison against NaN is false, so neither the low nor the high branch ever fires for it.

Gradient matches PyTorch’s clamp convention, not “zero at the boundary too”: the upstream gradient passes through unchanged where min <= x && x <= max – including exactly at either bound – and is zero everywhere x was actually moved by clamping. A NaN input therefore also gets a zero gradient, since x >= min is itself false for NaN under the same ordered-comparison rule.

pub fn hardtanh(&self, min_val: f32, max_val: f32) -> Result<Self>

Elementwise Hardtanh, PyTorch’s F.hardtanh(x, min_val, max_val): the same values as Self::clamp(Some(min_val), Some(max_val)), but with hardtanh’s OWN kink convention rather than clamp’s. PyTorch’s hardtanh_backward tests self <= min_val || self >= max_val (zero there, at BOTH bounds, not just outside them), so the gradient passes through only strictly inside (min_val, max_val) – unlike Self::clamp, whose own gradient is inclusive of both bounds. Composed from three disjoint region masks (<= min_val, the open interior, >= max_val) rather than reusing clamp’s Rule, so the boundary convention is exact even when min_val == max_val. Rejects a non-finite bound or min_val > max_val before any device launch, matching clamp.

pub fn relu6(&self) -> Result<Self>

Elementwise ReLU6, PyTorch’s F.hardtanh(x, 0., 6.): Self::hardtanh with fixed bounds 0 and 6, including its strictly-interior gradient at both kinks.

pub fn hardsigmoid(&self) -> Result<Self>

Elementwise Hardsigmoid, PyTorch’s piecewise-linear sigmoid approximation: clamp(x + 3, 0, 6) / 6. Backward matches PyTorch’s hardsigmoid_backward: grad / 6 strictly inside (-3, 3), exactly 0 at and outside either bound (self > -3 && self < 3, both strict). Composed from an interior mask and a saturated-high mask rather than clamp, for the same boundary reason as Self::hardtanh.

pub fn hardswish(&self) -> Result<Self>

Elementwise Hardswish, PyTorch’s x * relu6(x + 3) / 6. Backward matches PyTorch’s hardswish_backward exactly, including its asymmetric kinks: 0 for x <= -3, grad * (x / 3 + 0.5) strictly inside (-3, 3), and grad (pass-through, derivative 1) for x >= 3 – so, unlike Self::hardtanh/Self::hardsigmoid, the two kinks are NOT symmetric: x == -3 routes to the zero branch but x == 3 routes to the pass-through branch, not the interior formula (whose limit there is 1.5, not 1). The interior region is composed as x * (x + 3) / 6, whose own derivative x / 3 + 0.5 falls out of the product rule for free.

pub fn hardshrink(&self, lambd: f32) -> Result<Self>

Elementwise Hardshrink: 0 where |x| <= lambd, otherwise x (PyTorch’s torch.where(x.abs() <= lambd, 0, x)). Backward matches PyTorch’s shared shrink_backward_kernel (also behind Self::softshrink): 0 on the closed band -lambd <= x <= lambd (inclusive of both bounds), grad outside it. Composed as x * outside_mask, so this natural product-rule derivative already equals that target with no extra masking trick. Rejects a non-finite or negative lambd before any device launch.

pub fn softshrink(&self, lambd: f32) -> Result<Self>

Elementwise Softshrink: x - lambd where x > lambd, x + lambd where x < -lambd, otherwise 0 on the closed band -lambd <= x <= lambd (PyTorch’s torch.where(x.abs() > lambd, x - x.sign() * lambd, 0)). Backward shares Self::hardshrink’s exact shrink_backward_kernel convention: 0 on that same closed band, grad outside it. Composed from two directional shift terms (each (x -+ lambd) * mask) whose masks are already mutually exclusive and complementary to the zero band, so the natural derivative equals the target directly. Rejects a non-finite or negative lambd before any device launch.

pub fn threshold(&self, threshold: f32, value: f32) -> Result<Self>

Elementwise Threshold(threshold, value): x where x > threshold, otherwise the constant value (PyTorch’s torch.where(x <= threshold, value, x); note the <=/> split, shared by forward and backward alike). Backward matches PyTorch’s threshold_backward: grad where x > threshold, exactly 0 at and below it. Composed as x * mask + value * (1 - mask), so the natural product-rule derivative already equals that target with no extra masking trick. Rejects a non-finite threshold or value before any device launch.

pub fn softsign(&self) -> Result<Self>

Elementwise Softsign, PyTorch’s x / (1 + x.abs()). No true kink: although the forward composes through Self::abs, its derivative 1 / (1 + x.abs())^2 is continuous through x == 0 (both one-sided limits agree), and the ordinary composed chain rule already reproduces it exactly there because abs’s own backward is exactly zero at x == 0 (its documented convention) rather than a one-sided slope, which is exactly the missing term the analytic derivative also drops at that point.

pub fn tanhshrink(&self) -> Result<Self>

Elementwise Tanhshrink, PyTorch’s x - x.tanh(). Smooth everywhere; composed directly from Self::tanh, so the chain rule gives the exact derivative 1 - (1 - tanh(x)^2) == tanh(x)^2 for free, with no dedicated rule.

pub fn mean(&self, axes: impl IntoAxes) -> Result<Self>

pub fn sum(&self, axes: impl IntoAxes) -> Result<Self>

Reduce precisely the named axes to their sum; every other axis is preserved unchanged. Backward broadcasts the upstream gradient across the reduced axes without scaling it, since each contributing element has unit local derivative. This runs the same group-sum kernel mean uses, with the constant scale 1.0 in place of mean’s 1 / extent; the two are the same primitive, not a division after the fact.

pub fn moments(&self, axes: impl IntoAxes) -> Result<(Self, Self)>

Population mean and variance over exactly the declared named axes.

pub fn mean_square(&self, axes: impl IntoAxes) -> Result<Self>

Population mean of squares over exactly the declared named axes.

pub fn min(&self, axis: Axis) -> Result<Self>

Reduce one named axis to its minimum finite value.

Non-finite candidates are ignored. A group without a finite candidate returns NaN and has zero derivative even if its upstream derivative is non-finite. Ties route the derivative to the first logical coordinate along axis, independently of physical layout.

pub fn adaptive_avg_pool3d( &self, spatial: [Axis; 3], target: [usize; 3], ) -> Result<Self>

Reduce three named spatial axes to a fixed target extent each, using PyTorch’s adaptive average-pooling bin formula per axis: start = floor(i * input / output), end = ceil((i + 1) * input / output). Every other axis (such as batch or channel) is preserved. When an axis’s input extent does not divide evenly by its target extent, adjacent bins can share one boundary element; that element is averaged, at full weight, into each bin it falls in, exactly as F.adaptive_avg_pool3d computes it. Unlike a fixed kernel/stride pool, bins have no padding: every bin is a nonempty range of real input coordinates.

pub fn adaptive_avg_pool1d(&self, spatial: Axis, target: usize) -> Result<Self>

Same contract as Self::adaptive_avg_pool3d over one named spatial axis.

pub fn adaptive_avg_pool2d( &self, spatial: [Axis; 2], target: [usize; 2], ) -> Result<Self>

Same contract as Self::adaptive_avg_pool3d over two named spatial axes.

pub fn adaptive_max_pool3d( &self, spatial: [Axis; 3], target: [usize; 3], ) -> Result<Self>

Reduce three named spatial axes to a fixed target extent each, taking each of Self::adaptive_avg_pool3d’s own bins to its maximum instead of its weighted mean.

Composed from Self::gather, Self::merge and Self::max – no dedicated kernel. Each spatial axis in turn gathers every output bin’s (host-computed, data-independent) member positions onto a fresh per-axis patch axis, padding a bin shorter than the widest one by repeating its own last real position. The per-axis patch axes are then merged, in the supplied spatial order with the last axis fastest, into ONE window axis that a single Self::max removes. That merged axis visits the whole Cartesian window in exactly PyTorch’s adaptive_max_pool CPU scan order (nested loops, last spatial axis innermost, strict > so the first maximum met is kept), and Self::max keeps the first logical maximum, so a tie – including one spanning two or more axes – routes the gradient to the element PyTorch’s returned index names. A padding repeat always comes after the real position it copies in that scan, so it can only tie, never win. gather’s backward is an exact scatter-add, so a position shared by two adjacent bins receives a gradient contribution from every bin it wins. Non-finite candidates are ignored as Self::max ignores them.

pub fn adaptive_max_pool1d(&self, spatial: Axis, target: usize) -> Result<Self>

Same contract as Self::adaptive_max_pool3d over one named spatial axis.

pub fn adaptive_max_pool2d( &self, spatial: [Axis; 2], target: [usize; 2], ) -> Result<Self>

Same contract as Self::adaptive_max_pool3d over two named spatial axes.

pub fn max_unpool3d( &self, spatial: [Axis; 3], indices: &[usize], output: [usize; 3], ) -> Result<Self>

Scatter each element of self into a zero tensor whose three named spatial axes take the extents output, at the flat spatial position indices names: PyTorch’s F.max_unpool3d. indices holds one entry per element of self, in self’s logical (Self::to_vec) order, each a row-major offset into the output spatial volume (the supplied spatial order, last axis fastest) – the convention MaxPool3d’s forward_with_indices returns, i.e. PyTorch’s return_indices=True indices into the flattened D*H*W plane of each batch/channel slice. Every other axis carries through unchanged; a position no index names is an exact zero.

When two elements of one slice name the same position (possible after an overlapping max pool), the output holds the LAST such element in logical order, matching PyTorch’s sequential CPU kernel (output[index] = input[i], last write wins); a non-finite value at an overwritten (non-final) duplicate makes that output position NaN. The gradient with respect to self is PyTorch’s own: every element, overwritten or not, receives the upstream gradient at the position it names (grad_output.gather(indices)); the indices themselves carry none. Composed from Self::merge, Self::gather, Self::scatter_add and Self::detach, with no new kernel.

pub fn max_unpool1d( &self, spatial: Axis, indices: &[usize], output: usize, ) -> Result<Self>

Same contract as Self::max_unpool3d over one named spatial axis.

pub fn max_unpool2d( &self, spatial: [Axis; 2], indices: &[usize], output: [usize; 2], ) -> Result<Self>

Same contract as Self::max_unpool3d over two named spatial axes.

pub fn max(&self, axis: Axis) -> Result<Self>

Reduce one named axis to its maximum finite value.

Non-finite candidates are ignored. A group without a finite candidate returns NaN and has zero derivative even if its upstream derivative is non-finite. Ties route the derivative to the first logical coordinate along axis, independently of physical layout. Exact mirror of Self::min.

pub fn logsumexp(&self, axis: Axis) -> Result<Self>

Numerically stable log-sum-exp reduction over one named axis: ln(sum(exp(x))) computed as m + ln(sum(exp(x - m))) with m = self.max(axis), matching torch.logsumexp. m is Self::detached before the subtraction, so no gradient flows back through it; differentiating the remaining composition by hand collapses to exactly softmax(x) along axis (exp(x - m) / sum(exp(x - m))), independent of m’s own derivative – the standard stabilizing-shift trick. Composed entirely from Self::max, Self::sub, Self::exp, Self::sum, Self::ln, and Self::add – no dedicated kernel or backward rule – so even a row whose entries sit near a large shared magnitude (e.g. near 1000) never exponentiates anything larger than zero and cannot overflow. Shares max’s axis-removal contract and its no-finite-candidate convention: a group with every entry non-finite gives m = NaN, so logsumexp returns NaN there too, unlike PyTorch’s torch.logsumexp, which returns -infinity for an all--infinity group. Unlike max, this composition has no dedicated backward rule to zero that group’s gradient, so NaN also propagates through it ordinarily, rather than landing on zero.

pub fn masked_mean(&self, mask: &Self) -> Result<Self>

Reduce a tensor to the mean of elements selected by a constant binary mask. The mask must have the same named axes, and every axis must be reduced.

pub fn masked_softmax(&self, axis: Axis, mask: &Self) -> Result<Self>

Softmax over one named axis, with positions mask marks 0.0 receiving exactly zero probability and exactly zero gradient, and a group whose mask is entirely zero along axis returning all zeros for that group instead of NaN. mask is a constant {0.0, 1.0} tensor sharing every one of self’s named axes, validated exactly like Self::masked_mean’s mask (finite 0/1 values only, rejected if it requires gradients) – except an all-zero group along axis is the expected fully-masked case here, not an error the way an entirely empty mask is for masked_mean.

This is the named-axis equivalent of torch.softmax(score.masked_fill(~keep, -inf), dim) followed by torch.nan_to_num(..., neginf=0.0) for a fully-masked row – the exact pair every gated-attention consumer with a variable-length validity mask (InterfaceMIL, PhaseSeparableFusion) otherwise repeats by hand. Folding the nan_to_num cleanup into the op’s own contract, rather than leaving a NaN for the caller to catch, is a deliberate difference from PyTorch’s two-call composition.

Composed as softmax(self + (1 - mask) * MASKED_SOFTMAX_OFFSET) * any_valid, where any_valid (mask.max(axis)) is 1.0 for a group with any valid position and 0.0 for a fully-masked one, and MASKED_SOFTMAX_OFFSET is a large finite negative constant rather than literal -inf. An actually-infinite offset would make a fully-masked row’s own maximum -inf too, so the stable-softmax step computes exp(-inf - -inf), i.e. NaN, for every position in that row – and NaN * 0.0 is still NaN, so the trailing any_valid multiply could never clean it up. The finite offset keeps every intermediate value finite: a masked position’s probability underflows exp to exactly 0.0f32 once its row has any valid position (the offset sits far past f32’s roughly -104 underflow threshold at any realistic score magnitude, with wide margin below overflowing the addition itself to -inf), and a fully-masked row instead computes an ordinary finite (if meaningless) distribution that the trailing multiply by any_valid zeroes out cleanly, since it is never NaN.

Gradient into self at a masked position is exactly zero either way: softmax’s backward rule scales the upstream gradient by the (exactly-zero) probability there, so it does not matter that the offset’s own local derivative into self is nominally 1.0. A fully-masked group’s gradient is exactly zero throughout, for the same reason applied to the any_valid multiply: it zeros the upstream gradient into that group’s raw softmax before the softmax backward rule ever runs.

pub fn causal_mask(&self, query: Axis, key: Axis) -> Result<Self>

Mask key positions greater than query positions with -infinity. This is square, zero-offset self-attention; cached/offset attention is not supported.

pub fn prefix_causal_mask( &self, query: Axis, key: Axis, prefix: usize, ) -> Result<Self>

Causal mask whose first prefix key positions stay visible to every query. A block of memory or code tokens placed before the sequence is attended freely, and the remaining positions are strictly causal. This is a logical dense mask, not a sparse kernel; prefix == 0 is causal_mask.

pub fn softmax(&self, axis: Axis) -> Result<Self>

Numerically stable normalization along one named axis, preserving the logical shape. Rows must contain finite values or -infinity and at least one finite value.

pub fn unfold2d( &self, channels: Axis, spatial: [Axis; 2], patch: Dim, kernel: [usize; 2], ) -> Result<Self>

Extract valid stride-one 2D patches, replacing channels with one flattened patch axis. Reverse mode scatters and sums overlapping patch contributions into the input.

pub fn contract(&self, rhs: &Self, axes: impl IntoAxes) -> Result<Self>

Sum the named shared axes; align remaining shared axes, retain distinct ones.

pub fn outer(&self, rhs: &Self) -> Result<Self>

Explicit expansion: shared axes still align, but no axes are reduced.

pub fn rename(&self, from: Axis, to: Axis) -> Result<Self>

pub fn pad_zeros(&self, axis: Axis, before: usize, after: usize) -> Result<Self>

Add exact zeros before and after one named axis, preserving logical axis order. Zero padding on both sides shares the original storage. Otherwise the result is materialized on-device using rank-sized geometry, including for backward.

pub fn narrow(&self, axis: Axis, start: usize, length: usize) -> Result<Self>

Keep a contiguous, nonempty interval of one named axis without removing it. Reject out-of-range or overflowing intervals. A full-axis interval shares storage; other intervals use device-side copies and zero-scattered gradients.

pub fn select(&self, axis: Axis, coordinate: usize) -> Result<Self>

Select one logical coordinate of a named axis and remove that axis.

The compact backend computes offsets from rank-sized metadata rather than constructing one host index per tensor element.

pub fn argmin(&self, axis: Axis) -> Result<Vec<usize>>

Host-side index of the first logical coordinate achieving the minimum along one named axis, per remaining position, in the same order Self::min itself removes that axis.

This is a discrete evaluation-path primitive: it names which coordinate won a reduction, so unlike Self::min it carries no gradient of its own and stays entirely host-side. It never invents a second comparator: it reads back Self::min’s own device-computed minimum for each remaining position alongside this tensor’s raw storage, then rescans each group in ascending coordinate order for the first element bit-equal to that minimum – exactly min’s own tie rule (backend.rs’s grouped_minimum keeps the earlier candidate on a strict < comparison, so the first logical coordinate wins any tie) and its own finite-only comparison (a non-finite candidate never matches). A group with no finite candidate is an error: there is no coordinate a NaN result could name.

pub fn scatter_add( &self, axis: Axis, index: &[usize], bucket: Axis, bucket_count: usize, ) -> Result<Self>

Sum values into buckets named by a host-side integer label per position of axis, replacing axis with a new bucket axis of extent bucket_count. index[i] names the bucket position i of axis contributes to; every other axis carries through unchanged, and a bucket no position names is an exact zero, not an error, matching PyTorch’s Tensor.index_add_.

This is the exact transpose of Self::gather: gather’s forward picks one source row per output position, and its backward scatter-adds the upstream gradient back into every row that picked it. scatter_add’s forward is that scatter-add, built from the same host index array; its own backward is exactly gather’s forward pick over the identical index (reading each source position’s own bucket back out of the upstream gradient). The pairing needs no new backend kernel: it reuses Plan::gather/Plan::reverse’s CSR machinery the other way around from gather, keyed by this tensor’s own physical storage so reordered layouts read correctly. index is checked against axis’s extent and bucket_count before any device work, and shares Plan’s 16,777,216-contribution limit, counted against this tensor’s own size.

pub fn bincount( index: &[usize], bucket: Axis, bucket_count: usize, device: &Device, ) -> Result<Self>

Per-bucket counts of a host-side bucket-index array: the constant-ones case of Self::scatter_add, so hard_eval’s torch.bincount(lab, minlength=k) has a direct spelling. Adds no backend machinery of its own; a caller that already has a values tensor to scatter can divide its own scatter_add sum by this count instead of calling both.

pub fn gather(&self, axis: Axis, index: &[usize], output: Axis) -> Result<Self>

Gather rows of a named axis by an arbitrary host-computed integer index, replacing that axis with a new named axis laid out along the index.

index[i] selects one logical coordinate of axis for output position i of output; every other axis carries through unchanged and every coordinate is checked against axis’s extent before any device work. This is a row copy, not Embedding’s dense one-hot contraction, so it stays affordable at table sizes a one-hot vector could never reach: the index lives on the host as plain integers, never as device-side vocabulary weight. Backward scatter-adds the upstream gradient into the picked rows using the same host-built, CSR-grouped reduction plan mean and min already use for their own broadcast and reduction gradients (Plan::gather/Plan::reverse in runtime::backend), so a repeated index accumulates every contribution deterministically and a row that is never picked gets an exact zero gradient. It shares that plan’s 16,777,216-contribution limit, counted against the gathered output’s size (index.len() times the extent of every other axis), never against the table’s own row count.

pub fn broadcast_to(&self, shape: &Shape) -> Result<Self>

Broadcast self onto shape, an explicit target that must carry every one of self’s axes at self’s own extent; shape may add axes self lacks entirely (they read with stride zero) and may reorder self’s existing axes. This is the outer-broadcast primitive: elementwise add/sub/mul/div only ever align one operand’s axis set onto the other’s when it is already a subset (binary, algebra/tensor.rs), and reject two operands that each have an axis the other lacks as “cannot introduce an implicit outer product”. A caller with genuinely disjoint axis sets – a pixel-indexed operand and a site-indexed operand, say – broadcasts each one explicitly onto a shared [pixel, site, ...] shape first, then combines the results with an ordinary elementwise op: a.broadcast_to(&shape)?.sub(&b.broadcast_to(&shape)?) composes a torch.cdist-style outer pairwise difference from broadcast_to and sub alone, with no dedicated outer-product op. Backward sums the upstream gradient over every axis this call added, exactly as every other broadcast in the crate reduces an introduced axis (elementwise add/mul/div, binary_cross_entropy_with_logits_weighted’s pos_weight).

pub fn cosine_similarity( &self, rhs: &Self, axis: Axis, eps: f32, ) -> Result<Self>

Cosine similarity along one named feature axis: PyTorch’s CosineSimilarity, cos = (x1 . x2) / max(||x1||_2 * ||x2||_2, eps). Semantic fix (was wave 1’s #125, landed with a documented divergence): PyTorch’s own kernel (aten/src/ATen/native/Distance.cpp’s cosine_similarity) computes x1_norm = x1.mul(x1).sum(dim), x2_norm likewise, then at::clamp_min(x1_norm * x2_norm, eps * eps).sqrt() as the ONE shared denominator – a single joint clamp on the PRODUCT of squared norms, not two per-vector clamps each floored to eps individually. The previous Axis composition clamped ||x1|| and ||x2|| separately before multiplying, which is a different function whenever exactly one operand’s norm falls below eps (the joint form still lets that operand’s real, sub-eps norm shrink the denominator, since only the PRODUCT is floored). Every consumer’s existing PairwiseDistance-family test used inputs with both norms comfortably above eps, where the two formulas agree exactly, so no other row’s recorded output changes. eps matches PyTorch’s default 1e-8 and must be finite and positive. self/rhs require identical axis sets including axis; every other axis is preserved unreduced.

pub fn pairwise_distance( &self, rhs: &Self, axis: Axis, p: f32, eps: f32, ) -> Result<Self>

p-norm pairwise distance along one named feature axis: PyTorch’s PairwiseDistance, norm(self - rhs + eps, p). eps (PyTorch’s default 1e-6) is added to the raw difference BEFORE the norm – exactly where PyTorch’s own F.pairwise_distance adds it (aten/src/ATen/native/Distance.cpp’s pairwise_distance computes at::norm(x1 - x2 + eps, p, ...)), not as a denominator floor. keepdim=False is automatic: axis is removed like every other Axis reduction. p (PyTorch’s default 2.0) must be a positive real or f32::INFINITY (PyTorch’s own p-norm domain, torch.norm’s ord). p == 2.0 keeps the exact ((self - rhs + eps)^2).sum(axis).sqrt() composition wave 1 shipped, bit-exact; p == 1.0 is abs(self - rhs + eps).sum(axis); p == f32::INFINITY is abs(self - rhs + eps).max(axis) (PyTorch’s p -> inf limit); every other p is the general (sum(abs(diff)^p))^(1/p), composed from ln/exp since there is no dedicated elementwise power op (x^p = exp(p * ln(x)) for x >= 0; ln(0) = -inf and exp(-inf) = 0 give the mathematically correct zero contribution at an exactly-zero coordinate, matching f32/PyTorch’s own IEEE ln/exp at that boundary). self/rhs require identical axis sets including axis.

pub fn margin_ranking_loss( &self, rhs: &Self, target: &Self, margin: f32, ) -> Result<Self>

PyTorch’s MarginRankingLoss: max(0, -target * (self - rhs) + margin), elementwise and unreduced. target is a constant 1.0/-1.0 per element (checked before launch) and carries no gradient, matching categorical_cross_entropy_with_logits’s constant targets. margin (PyTorch default 0.0) must be finite. self, rhs, and target require identical axis sets.

pub fn hinge_embedding_loss(&self, target: &Self, margin: f32) -> Result<Self>

PyTorch’s HingeEmbeddingLoss: per element, self where target == 1, else max(0, margin - self) where target == -1. target is a constant 1.0/-1.0 per element (checked before launch, no gradient). margin (PyTorch default 1.0) must be finite. self and target require identical axis sets. Unreduced.

pub fn cosine_embedding_loss( &self, rhs: &Self, target: &Self, feature: Axis, margin: f32, ) -> Result<Self>

PyTorch’s CosineEmbeddingLoss: 1 - cos(self, rhs) where target == 1, else max(0, cos(self, rhs) - margin) where target == -1, reducing the named feature axis via Self::cosine_similarity at its own default eps = 1e-8. PyTorch’s public documented formula names no eps at all; its C++ kernel instead adds a fixed, undocumented 1e-12 inside the sum of squares. Reusing CosineSimilarity’s own public default keeps one canonical epsilon across the family rather than inventing a second undocumented constant. margin (PyTorch default 0.0) must be finite. target is a constant 1.0/-1.0 (checked before launch): it must not contain feature, and its remaining axes must be a subset of self’s. Unreduced over every axis but feature.

pub fn triplet_margin_loss( &self, positive: &Self, negative: &Self, axis: Axis, margin: f32, p: f32, eps: f32, swap: bool, ) -> Result<Self>

PyTorch’s TripletMarginLoss: max(0, margin + d(self, positive) - d(self, negative)), where d is Self::pairwise_distance along axis at the given p (PyTorch default 2.0, any positive real or f32::INFINITY, exactly pairwise_distance’s own domain) and eps. When swap is true (PyTorch default false), d(self, negative) is instead the smaller of itself and d(positive, negative) (Balntas et al.’s swap term, penalizing a positive that sits closer to the negative than the anchor does), computed by stacking the two distances on a fresh axis and reducing with Self::min. margin (PyTorch default 1.0) must be finite. Unreduced over every axis but axis.

pub fn triplet_margin_with_distance_loss( &self, positive: &Self, negative: &Self, margin: f32, swap: bool, distance: impl Fn(&Self, &Self) -> Result<Self>, ) -> Result<Self>

PyTorch’s TripletMarginWithDistanceLoss: exactly Self::triplet_margin_loss’s max(0, margin + d(self, positive) - d(self, negative)) and swap rule, with d an arbitrary caller-supplied Rust closure over two tensors in place of a fixed p-norm. PyTorch defaults distance_function to None, meaning PairwiseDistance(); Rust has no Option-shaped default that keeps a plain, statically dispatched closure parameter, so the default is spelled explicitly by the caller, e.g. anchor.triplet_margin_with_distance_loss(&pos, &neg, margin, swap, |a, b| a.pairwise_distance(b, axis, 2.0, eps))?. margin (PyTorch default 1.0) must be finite.

pub fn multi_margin_loss( &self, target: &Self, class: Axis, p: f32, margin: f32, weight: Option<&Self>, ) -> Result<Self>

PyTorch’s MultiMarginLoss: for each class i, w[target] * max(0, margin - self[target] + self[i])^p, summed over i != target and scaled by 1 / class.extent(). target is a constant one-hot indicator over class (checked before launch: every value 0.0/1.0, summing to exactly 1.0), matching how categorical_cross_entropy_with_logits already spells a class label as a tensor rather than a host index. p (PyTorch default 1.0) must be exactly 1.0 or 2.0, PyTorch’s own documented domain (p == 1 or p == 2 are the only values its forward formula supports). weight (PyTorch default None, meaning every class weighs 1.0) is an optional length-class.extent() per-class tensor indexed by class alone; when given, PyTorch’s kernel (aten/src/ATen/native/LossMulti.h’s multi_margin_loss) multiplies the WHOLE per-sample hinge sum by weight[target] – the true class’s own weight, gathered once per sample via target.mul(weight).sum(class) – never a weight averaged or looked up per competing class i. margin (PyTorch default 1.0) must be finite; class must have at least two classes. Unreduced over every axis but class.

pub fn multi_label_margin_loss( &self, target: &Self, class: Axis, ) -> Result<Self>

PyTorch’s MultiLabelMarginLoss: (1 / class.extent()) * sum_{i,j} max(0, 1 - (self[j] - self[i])), summed over class positions i that are NOT a positive label and j that ARE. PyTorch spells the positive-label set as a fixed-width index array terminated by -1; Axis instead takes target as a constant multi-hot {0.0, 1.0} indicator over class (1.0 at every positive label) – the same floating-point-tensor spelling categorical_cross_entropy_with_logits already uses for a single label, generalized to a set. The two encodings name the same label sets; target is checked before launch (every value exactly 0.0 or 1.0) and carries no gradient. class must have at least two classes. Neither margin (PyTorch fixes it at 1.0; it is not a parameter of this class) nor weight/reduction exist to configure. Unreduced over every axis but class.

pub fn stack(values: &[Self], axis: Axis, position: usize) -> Result<Self>

Stack equal named shapes, inserting a new logical axis at position. Physical storage is stack-major so each source remains one contiguous copy.

pub fn concat(values: &[Self], axis: Axis) -> Result<Self>

Concatenate two or more tensors along a named axis they already share, extending that axis’s extent by the sum of each operand’s extent. Every other axis must match exactly across operands (same identity, same extent); the output’s axis order follows the first operand’s. This is the wave-1-proven composition of pad_zeros and add: each operand is zero-padded into its own slice of the concatenated axis, then the padded tensors are summed, so backward automatically narrows the incoming gradient back to each operand’s slice with no dedicated rule.

pub fn roll(&self, axis: Axis, shift: isize) -> Result<Self>

Cyclically shift one named axis’s values by shift, wrapping values that run off one end onto the other. PyTorch’s torch.roll sign convention: the element at logical index i moves to (i + shift).rem_euclid(extent), so a positive shift moves values toward higher indices and a negative shift moves them toward lower indices. Shifting by 0 or by any multiple of the axis’s extent is the identity. Rolling one axis at a time composes exactly with PyTorch’s multi-axis dims=: torch.roll(x, shifts=(-s, -s), dims=(1, 2)) is two calls, one per axis (SwinIR’s shifted-window attention, network_swinir.py:251,271).

Implemented as concat of two narrow slices – the wrapped tail moved ahead of the untouched head – rather than a dedicated backend rule: both primitives already carry exact device kernels and exact gradients, so backward falls out for free as the inverse roll (the same call with shift negated), with no new Rule variant to maintain.

pub fn with_layout(&self, order: impl IntoAxes) -> Result<Self>

Materialize a storage order without changing logical axes or values.

pub fn split( &self, axis: Axis, dims: impl IntoIterator<Item = Dim>, ) -> Result<Self>

pub fn merge(&self, axes: impl IntoAxes, axis: Axis) -> Result<Self>

Selected-axis order defines flattening; insert the merged axis at the first selected position.

pub fn upsample_nearest(&self, axis: Axis, factor: usize) -> Result<Self>

F.interpolate(mode="nearest") for one named spatial axis at an exact integer scale factor. Composed entirely from stack (duplicate the input factor times along a fresh axis) and merge (fold that axis into the named spatial axis, spatial-axis major, so each source element becomes factor adjacent identical copies). This is the composition wave 1 of the research migration proved bit-exact against a real F.interpolate(mode="nearest") oracle (research/src/vision/morpheus/axis/tests/partitioner_trunk_fpn_obj.rs, upsample2x, filed as axis issue #66). morpheus‘s mobilesam/scale10/train_partitioner.py:107 and detector/train_detector.py:237,239 both call this at factor = 2 on height then width separately (stride-16 to stride-8, stride-8 to stride-4 in an FPN fusion); call this once per spatial axis to reproduce that. Backward needs no dedicated rule: stack and merge are already differentiable, so the gradient for a source element is the sum of its factor output copies’ incoming gradients.

pub fn resample_bilinear(&self, axis: Axis, out_extent: usize) -> Result<Self>

torch.nn.Upsample(mode="bilinear", align_corners=False) for one named axis, at any positive output extent (world/fluid’s scripts/train.py:101-102 only ever calls this at scale_factor=2, i.e. out_extent = 2 * self.extent(axis), applied to height then width separately before each of the U-Net decoder’s two skip concatenations; filed as axis issue #66). The interpolation weights are fixed by the input/output extents alone (never by tensor data), so this is one contract against a host-built [axis, resampled] weight matrix rather than a dedicated kernel; contract’s existing backward is already the transpose of that same matrix, so the gradient is exact for free and distributes to up to two source coordinates per output coordinate (up to four source cells when a height and a width axis are each resampled, one call per axis).

Weights follow PyTorch’s own align_corners=False half-pixel formula: for output coordinate j, source = max(0, (j + 0.5) * in_extent / out_extent - 0.5); let lower = floor(source) clamped to the last valid input coordinate and upper = lower + 1 (or lower again at the last input coordinate, so its weight is exactly 1). lower receives weight 1 - fract(source) and upper receives weight fract(source).

pub fn resample_bicubic(&self, axis: Axis, out_extent: usize) -> Result<Self>

torch.nn.functional.interpolate(mode="bicubic", align_corners=False) for one named axis, at any positive output extent. Closes part of Upsample’s docs/nn/catalog.md “partial” gap (“no trilinear or bicubic”). Same construction as Self::resample_bilinear: weights are fixed by the input/output extents alone, so this is one Self::contract against a host-built [axis, resampled] weight matrix, not a dedicated kernel, and inherits contract’s exact transpose backward for free.

Weights use PyTorch’s own separable bicubic convolution kernel with a = -0.75 (torch.nn.functional.interpolate’s literal constant): W(x) = (a+2)|x|^3 - (a+3)|x|^2 + 1 for |x| <= 1, W(x) = a|x|^3 - 5a|x|^2 + 8a|x| - 4a for 1 < |x| < 2, else 0. Source coordinates use the same align_corners=False half-pixel formula as resample_bilinear (source = (j + 0.5) * in/out - 0.5, unclamped here since the kernel itself decays to 0 past its support); each output column draws from the four taps floor(source) - 1 ..= floor(source) + 2, each tap coordinate clamped (border-replicated) into [0, in_extent - 1] exactly as PyTorch’s CPU/CUDA kernels do, and taps that clamp to the same input coordinate have their weights summed rather than overwritten. align_corners=True is not implemented: out of scope for the row’s stated gap, which names only bicubic, trilinear (already proven via composed resample_bilinear) and scale_factor=.

pub fn backward(&self) -> Result<()>

Reverse mode from a scalar. Releases the graph after success; rebuild it for another backward.

Trait Implementations

impl Clone for Tensor

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.