Easy Primitives

Implement Cosine Similarity

Cosine similarity compares the directions of two nonzero vectors. It is their dot product divided by the product of their lengths:

similarity = sum(a * b) / (sqrt(sum(a * a)) * sqrt(sum(b * b)))

For a = [1, 2] and b = [2, 1], the dot product is 4 and each length is sqrt(5). Their similarity is therefore 4 / 5 = 0.8.

In exact arithmetic, the result lies between -1 and 1. A value of 1 means the vectors point in the same direction, 0 means they are perpendicular, and -1 means they point in opposite directions. Positive rescaling does not change the result: [1, 0] and [10, 0] have similarity 1.

Your task

Implement cosine_similarity(a, b) for two nonempty 1D floating-point tensors of the same shape. Return a scalar tensor using the formula above. Use basic tensor operations rather than a built-in cosine-similarity function.

Both vectors have nonzero length. Zero vectors have no direction and are outside this exercise's input contract; you do not need an epsilon rule. Tests use values whose squared sums and denominator are representable. Allow normal floating-point rounding in the result.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Learn vectors that express similarity

An embedding maps an example to a vector. Metric learning trains these vectors so that related examples are close or similar under a chosen comparison. Cosine similarity compares direction:

$$\operatorname{sim}(a,b)=\frac{a\cdot b}{\|a\|\|b\|}.$$

Vectors [1, 0] and [2, 0] have similarity 1 despite different lengths. Vectors [1, 0] and [0, 1] have similarity 0. A zero vector has no direction, so follow the task's rule for a zero denominator or epsilon.

Compare an anchor, a positive, and a negative

Triplet loss uses an anchor a, a related positive p, and an unrelated negative n. A common formula is

$$\max(0,d(a,p)-d(a,n)+m).$$

The margin m specifies how much farther away the negative should be. For positive distance 2, negative distance 4, and margin 1, the loss is zero. If the negative distance is 2.5, the loss is 0.5.

Euclidean distance and squared Euclidean distance give different losses and gradients. Use the definition in the exercise.

Select the positive from several candidates

InfoNCE treats the matching task as classification. Divide similarity scores by a positive temperature to get logits, then apply cross-entropy with the positive candidate as the target. A lower temperature concentrates probability more strongly on the largest scores without changing their order.

Check which axis contains candidates, which index is the positive, and whether to normalize embeddings first. Paired batches often place positives on the diagonal, but only when their ordering supports that convention. Computing a loss in one matching direction differs from averaging losses in both directions.

Detecting runtime…

Implement Cosine Similarity

Your task

Learn vectors that express similarity

Compare an anchor, a positive, and a negative

Select the positive from several candidates

Your code, saved here

Stuck?

JAX reference solution