Debug a PyTorch Contrastive Loss

You are given a PyTorch function intended to compute a contrastive loss. The function takes two tensors, z1 and z2, each of shape [batch_size, embedding_dim], and a temperature parameter. Row i in z1 is considered the positive partner of row i in z2. All other rows in z2 are considered negative examples for the anchor in row i of z1.

The current implementation attempts to compute a score matrix using torch.matmul(z1, z2) and then applies F.cross_entropy with labels derived from the batch indices.

Your task is to:

1. Explain why the current implementation fails to compute the intended contrastive loss, focusing on the matrix multiplication shape mismatch and the resulting score matrix.
2. Repair the score matrix calculation to correctly represent the similarity scores between all z1 anchors and all z2 rows.
3. Validate the function's assumptions about the input tensors and parameters.
4. Discuss and implement strategies for handling numerical precision, particularly when using lower-precision floating-point types (e.g., float16).

The primary goal is to compute the mean cross-entropy for z1 anchors scored against all z2 rows. An optional extension is to implement a symmetric average with the reverse direction. L2 normalization should only be added if the similarity is explicitly defined as cosine similarity.

Constraints & Assumptions

• Both z1 and z2 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 (i.e., z1[i] is the positive for z2[i]).
• For float16 or bfloat16 inputs, logits and loss computation should occur 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.

Similar LeetCode problem: Dot Product of Two Sparse Vectors