Correct p-norm computation in triplet_loss

[pchintar]

Description

nn.losses.triplet_loss exposes a configurable norm degree through the p argument and documents the distance term as a p-norm:

||A - P||_p

However, the current distance computation is:

mx.sqrt(mx.power(anchors - positives, p).sum(axis) + eps)

which only matches the documented p-norm when p=2. For other values of p, it can produce incorrect results and even NaN values.

This change updates the distance computation to use the standard p-norm:

(sum(abs(x - y) ** p) + eps) ** (1 / p)

which matches both the function documentation and MLX's linalg.norm behavior.

On a sidenote, I also noticed that the LayerNorm, GroupNorm, and BatchNorm docstrings in normalization.py had epsilon outside the square root (\sqrt{Var[x]} + \epsilon) which doesn't match the standard normalization formula, so I changed it to (\sqrt{Var[x] + \epsilon}) instead.

Changes

• Corrected the p-norm computation in nn.losses.triplet_loss
• Added test coverage for non-default p values

Testing

python -m unittest python.tests.test_losses.TestLosses.test_triplet_loss

cc @zcbenz

[zcbenz]

Looks good to me, thanks for fixing this!

[pchintar]

@zcbenz thnx again