torchcell.losses.multi_dim_nan_tolerant.SupCR

class SupCR(temperature: float = 0.1, eps: float = 1e-07)

Bases: Module

Supervised Contrastive Regression (SupCR) loss over multiple target dimensions.

Uses efficient O(N^2 log N) implementation with row-wise sorting and suffix sums.

compute_similarity(embeddings: Tensor) → Tensor

Compute cosine similarities between embeddings.

compute_dimension_loss(embeddings: Tensor, labels: Tensor) → Tensor

Compute SupCR loss for a single dimension.

Parameters:

• embeddings – Shape [batch_size, embedding_dim]

• labels – Shape [batch_size] for single dimension

forward(embeddings: Tensor, labels: Tensor) → Tensor

Compute losses for all dimensions.

Parameters:

• embeddings – Shape [batch_size, embedding_dim]

• labels – Shape [batch_size, num_dims]