Debug a PyTorch Contrastive Loss

The source reports a PyTorch contrastive-loss debugging interview centered on matrix-multiplication shape mismatches. It does not provide the exact code or report a normalization-axis defect. The snippet and one-directional in-batch objective below are a practice reconstruction of the reported shape bug.

import torch
import torch.nn.functional as F

def contrastive_loss(z1, z2, temperature):
 logits = torch.matmul(z1, z2) / temperature
 labels = torch.arange(z1.shape[0], device=z1.device)
 return F.cross_entropy(logits, labels)

z1 and z2 each have shape [batch_size, embedding_dim]. Row i in z1 is the positive partner of row i in z2; the other rows of z2 are negatives for that anchor. Explain the failed contraction, repair the score matrix, validate the function's assumptions, and discuss numerical precision.

The required output is the mean cross-entropy for z1 anchors scored against all z2 rows. A symmetric average with the reverse direction is an optional extension, not part of the required repair. Likewise, L2 normalization is a modeling choice: add it only if the similarity is explicitly defined as cosine similarity.

Constraints & Assumptions

• Both inputs are finite rank-two floating tensors on the same device, with the same shape and dtype.
• batch_size > 0 and embedding_dim > 0.
• temperature is a finite positive Python int or float and is not a Boolean.
• Positive pairs are aligned by row index.
• For float16 or bfloat16 inputs, form logits and compute loss in float32 while preserving gradient flow. Float32 and float64 inputs may retain their dtype.
• Raise ValueError for invalid inputs and FloatingPointError if the resulting logits are non-finite.

Clarifying Questions to Ask

• Is similarity a raw dot product or cosine similarity?
• Are negatives limited to the opposite view?
• Can one anchor have multiple positives?
• Is a one-directional or symmetric objective intended?
• What mixed-precision policy is expected?