torchcell.losses.multi_dim_nan_tolerant.WeightedDistLoss

class WeightedDistLoss(bandwidth: float = 0.5, weights: Tensor | None = None, regularization_strength: float = 0.1, loss_fn: str = 'L2', eps: float = 1e-07)

Bases: Module

Distribution-matching (DistLoss) loss using KDE and differentiable sorting.

forward(y_pred: Tensor, y_true: Tensor) → tuple[Tensor, Tensor]

Compute distribution loss for multi-dimensional outputs.

Parameters:

• y_pred – Predictions [batch_size, num_dims]

• y_true – True labels [batch_size, num_dims]

Returns:

Weighted sum of dimension losses dim_losses: Individual losses per dimension

Return type:

weighted_loss