torchcell.losses.multi_dim_nan_tolerant.CategoricalEntropyRegLoss

class CategoricalEntropyRegLoss(lambda_d: float = 0.1, lambda_t: float = 0.1, num_classes: int = 32)

Bases: Module

Entropy regularization loss that works with both soft and one-hot categorical targets.

compute_pairwise_distances(x: Tensor) → Tensor

Compute pairwise L2 distances between all points.

compute_target_distances(targets: Tensor) → Tensor

Compute KL divergence distances between probability distributions.

compute_centers(features: Tensor, targets: Tensor) → dict[int, dict[int, Tensor]]

Compute weighted feature centers for each class in each dimension.

forward(features: Tensor, targets: Tensor, mask: Tensor) → tuple[Tensor, Tensor, Tensor]

Forward pass computing diversity and tightness losses.

Parameters:

• features – Node features [batch_size, feature_dim]

• targets – Probability distributions [batch_size, num_dims * num_classes]

• mask – Valid sample mask [batch_size]