Loss

TorchLean loss helpers in the style of torch.nn.functional.

These helpers keep training loops close to the familiar torch.nn.functional style:

loss = Loss.mse yhat y
loss = Loss.mse yhat y (reduction := .sum)

They are execution-mode generic: eager tape and typed SSA/DAG both work.

PyTorch references

• torch.nn.functional (losses overview): https://pytorch.org/docs/stable/nn.functional.html
• mse_loss: https://pytorch.org/docs/stable/generated/torch.nn.functional.mse_loss.html
• cross_entropy: https://pytorch.org/docs/stable/generated/torch.nn.functional.cross_entropy.html
• nll_loss: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html
• binary_cross_entropy_with_logits: https://pytorch.org/docs/stable/generated/torch.nn.functional.binary_cross_entropy_with_logits.html

inductive TorchLean.Loss.Reduction : Type

Reduction mode for losses that start as elementwise tensors.

PyTorch analogy: reduction="mean" or reduction="sum".

• mean : Reduction
• sum : Reduction

@[instance_reducible]

instance TorchLean.Loss.instReprReduction : Repr Reduction

def TorchLean.Loss.instReprReduction.repr : Reduction → ℕ → Std.Format

@[instance_reducible]

instance TorchLean.Loss.instDecidableEqReduction : DecidableEq Reduction

def TorchLean.Loss.reduceLoss {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {s : Shape} (x : Runtime.Autograd.Model.RefTy m α s) (reduction : Reduction) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Reduce an elementwise loss tensor to a scalar according to reduction.

This is the common final step for losses like MSE and cross-entropy.

def TorchLean.Loss.mse {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {s : Shape} (yhat y : Runtime.Autograd.Model.RefTy m α s) (reduction : Reduction := Reduction.mean) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Mean squared error (MSE) loss between predictions and targets.

This is backend-generic and supports both mean and sum reduction.

def TorchLean.Loss.mseWeighted {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {s : Shape} (yhat y weights : Runtime.Autograd.Model.RefTy m α s) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Weighted mean-squared error.

This returns sum (weights * (prediction - target)^2) without implicit normalization. Weights whose sum is one therefore define a weighted mean, and zero weights exclude coordinates without requiring a separate masking operation.

def TorchLean.Loss.nllOneHot {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logProbs targetOneHot : Runtime.Autograd.Model.RefTy m α s) (reduction : Reduction := Reduction.mean) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Negative log-likelihood (one-hot targets), assuming inputs are log-probabilities.

def TorchLean.Loss.oneHotCrossEntropy {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {s : Shape} (axis : ℕ) [Shape.AxisInBounds axis s] (logits targetOneHot : Runtime.Autograd.Model.RefTy m α s) (reduction : Reduction := Reduction.mean) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Cross-entropy (one-hot targets), computed as -sum(y * log(softmax(logits))).

Note: this is one-hot only. For integer class labels, use crossEntropy, whose Fin classes target tensor records the class bound in its element type.

def TorchLean.Loss.nllUnreduced {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {leading trailing : Shape} {classes : ℕ} (classAxis : ℕ) (_hClassAxis : classAxis = leading.rank) (logProbs : Runtime.Autograd.Model.RefTy m α (leading.concat (trailing.prependDim classes))) (target : Runtime.Autograd.Torch.DataRef (Fin classes) (leading.concat trailing)) : m (Runtime.Autograd.Model.RefTy m α (leading.concat trailing))

Unreduced indexed negative log-likelihood along an arbitrary class axis.

The shape contract is:

• logProbs has shape prefix ++ [classes] ++ suffix;
• classAxis = rank prefix identifies the classes dimension;
• target has shape prefix ++ suffix and element type Fin classes;
• the result has the same shape as target, with the class axis removed.

Each output coordinate is -logProbs[..., target[...], ...]. In particular, suffix coordinates remain aligned with their corresponding labels. Invalid class labels are unrepresentable.

def TorchLean.Loss.nll {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {leading trailing : Shape} {classes : ℕ} (classAxis : ℕ) (hClassAxis : classAxis = leading.rank) (logProbs : Runtime.Autograd.Model.RefTy m α (leading.concat (trailing.prependDim classes))) (target : Runtime.Autograd.Torch.DataRef (Fin classes) (leading.concat trailing)) (reduction : Reduction := Reduction.mean) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Indexed negative log-likelihood along an arbitrary class axis.

The logits, labels, and class-axis contract are those of nllUnreduced. The requested reduction is applied over every coordinate of the class-erased output shape.

def TorchLean.Loss.nllWeighted {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {leading trailing : Shape} {classes : ℕ} (classAxis : ℕ) (hClassAxis : classAxis = leading.rank) (logProbs : Runtime.Autograd.Model.RefTy m α (leading.concat (trailing.prependDim classes))) (target : Runtime.Autograd.Torch.DataRef (Fin classes) (leading.concat trailing)) (weights : Runtime.Autograd.Model.RefTy m α (leading.concat trailing)) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Weighted indexed negative log-likelihood along an arbitrary class axis.

The logits, labels, weights, and output coordinates obey the nllUnreduced shape contract. This returns sum (weights * losses) without implicit normalization; normalized weights therefore give a weighted mean, while zero weights mask coordinates without a division-by-zero convention.

def TorchLean.Loss.crossEntropy {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {leading trailing : Shape} {classes : ℕ} (classAxis : ℕ) (hClassAxis : classAxis = leading.rank) (logits : Runtime.Autograd.Model.RefTy m α (leading.concat (trailing.prependDim classes))) (target : Runtime.Autograd.Torch.DataRef (Fin classes) (leading.concat trailing)) (reduction : Reduction := Reduction.mean) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Indexed cross-entropy along an arbitrary class axis.

This applies log-softmax along classAxis and then uses nll. Shapes follow the prefix ++ [classes] ++ suffix contract documented on nllUnreduced.

def TorchLean.Loss.crossEntropyWeighted {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {leading trailing : Shape} {classes : ℕ} (classAxis : ℕ) (hClassAxis : classAxis = leading.rank) (logits : Runtime.Autograd.Model.RefTy m α (leading.concat (trailing.prependDim classes))) (target : Runtime.Autograd.Torch.DataRef (Fin classes) (leading.concat trailing)) (weights : Runtime.Autograd.Model.RefTy m α (leading.concat trailing)) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Weighted indexed cross-entropy along an arbitrary class axis.

This applies log-softmax along classAxis, multiplies each class-erased loss by the corresponding weight, and sums without implicit normalization. Shapes follow the contract on nllUnreduced.

def TorchLean.Loss.bceWithLogits {α : Type} [Storage α] [Context α] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] {s : Shape} (logits target : Runtime.Autograd.Model.RefTy m α s) (reduction : Reduction := Reduction.mean) : m (Runtime.Autograd.Model.RefTy m α Spec.Shape.scalar)

Binary cross-entropy with logits (elementwise), using the stable identity: BCEWithLogits(x,y) = y * softplus(-x) + (1-y) * softplus(x).

Targets are expected in [0,1] (typically 0/1), same shape as logits.