Add fixed size unique operation

[louen]

This PR adds a unique operator, as a convenience to match with similar operations found in numpy, torch and jax.

The operation is implemented as a composite of other primitive ops (sort and scatter) and is mostly useful as a parity convenience operation.

Size-dependency

One of the main challenges of a unique operation in MLX is that the size of the unique array is unknown before evaluation, which doesn't really fit MLX's design (see discussion in e.g. #246, #568 and #927, decision in @awni's gist)

This issue is also a problem for jax and its jit compiler ; and this PR follows their approach to resolving this issue: a user-passed size parameter which determines the size of the output array independently from the input.

In the case of jax they propose the size-independent version as an option. The MLX implementation in this PR makes this a requirement.

Behaviours and edge-cases

size

The new function aims for parity of behaviour with the jax version (with the provided size) wherever possible.

The choice of size will determine what happens to the output array:

• if the provided size is smaller than the number of unique elements, the output will be truncated
• if the provided size is larger than the number of unique elements, the output will be padded with fill_value
• if a fill_value is not provided, the first element of the sorted array is used (this is usually the min value of the array, except in pathological cases).

An empty array is generally a valid input, but unique([], size) will throw if size is greater than 0 (since there is no min value to fill the result with). However, unique([], size, fill_value=-1) will return [-1] * size.

A size of 0 is supported, returning an empty array.

Optional return arrays

unique optionally returns the counts of each unique element and the inverse indices if requested.

If the unique array is padded, the counts for the padding elements are 0. This means that count.sum() == a.size() if and only if the output was not truncated. In the truncated case, the counts for the unique elements are truncated too, so count.sum() will be smaller.

The inverse array holds indices in unique that allow to reconstruct the original array. If the output is truncated, those indices are clamped at size-1, so that a = unique[inverse] should always be valid if size > 0, but will only reconstruct the original array if the unique array was not truncated. Otherwise, the bigger elements will be replaced by the largest kept value.

Note that inverse's dtype is uint32 (matching other mlx ops returning indices such as argmax) where jax uses int32.

Differentiation

Gradient propagates up the unique array back to a, deterministically through the first appearing representative of each of the unique array elements.

NaN handling

NaN values are handled and placed at the end of the sorted unique array (before padding); but consistently with NaN != NaN they do not collapse into one single unique value.
This is a significant disparity with jax (which collapses NaNs by default, but allows to opt out).
We could support NaN collapse at the cost of an extra argument and an extra check on the edge detection to match with jax's behaviour.

Limitations

• bool arrays are unsupported on metal (this is due to the fact that sort fails on bool arrays). Unique on a bool array is only guaranteed to work on CPU stream, at least on Metal platforms.
• complex arrays containing NaN can fail to sort properly on CPU stream, but this is a bug in complex sort ([BUG] argsort diverges between cpu and gpu on complexes with NaN #4502)
• axis, equal_nan and return_index optional parameters found in jax (et al.) are not supported yet.
• backwards gradient is supported, but jvp is not supported (unlike Jax). This is a limitation of current mlx's scatter.

AI usage disclosure:

• Claude code (Opus 5) ported and reimplemented my original custom implementation, reviewed, suggested improvements (differentiability, edge cases, comments, Jax parity checks), and wrote the python tests.
• PR messages and commit messages 100% written by me, a human

[RohanGautam]

Thanks for the PR!

[RohanGautam]

We should have the tests here include the variations in mx.unique - the case with return_inverse=True, and fill_value specified

[louen]

do you mean specifically for gradient computation or are you asking more generally for these parameters to be exercised more in the test ?

[louen]

Adding return_index is actually two lines so that seems like a low-hanging fruit.
There are a bunch of tests who exercise fill_value, but I can add some if you want.

One place I wasn't sure is when using a fill value that cannot be represented with the current type (e.g. using -1 on uint8. Jax accepts it but the current code with throw instead, which is what pad does too. Here I went for consistency over Jax parity.

[louen]

Now supports return_index as requested.

[zcbenz]

Thanks!