SIGN IN SIGN UP

Match numpy output shapes in jnp.unique for empty arrays

jnp.unique(x, axis=k) fabricated a phantom slice when the input was
empty along every dimension: for a (0, 0) array it returned shape
(1, 0) with axis=0 where numpy returns (0, 0), (0, 1) with axis=1,
and (4, 1, 0) or (4, 0, 1) for (4, 0, 0) inputs.

_unique_sorted_mask forces size=1 when math.prod(out_shape) == 0
because lexsort cannot sort zero keys of nonzero length; collapsing
existing empty slices to one also matches numpy, which treats them as
vacuously equal ((n, 0) inputs give (1, 0)). The collapse was applied
even when no slice exists at all (size == 0), fabricating a row numpy
does not produce. Clamp with min(size, 1): the lexsort workaround is
still avoided and nonempty slices still collapse, while a fully empty
axis now flows through the existing size-0 handling in _unique.

For fully-empty input with a concrete size, value padding is unchanged,
but return_index shape changes from (1,) to (0,), matching numpy.

setdiff1d has a separate zero-size issue rooted in _in1d, not this
helper, and is unchanged here.

Test Plan:
python -m pytest tests/lax_numpy_setops_test.py -n auto -q
  152 passed
JAX_ENABLE_X64=1 python -m pytest tests/lax_numpy_setops_test.py -n auto -q
  156 passed
Coverage for (0, 0) inputs flows through the shared empty_array_shapes
list at the top of the test file, which now contains (0, 0) and feeds
testUnique and the isin/setdiff1d/union1d sweeps.
V
Vaggelis committed
2d4a00d218671aaca15e66cf2c3ff6b3397a4c02
Parent: ccdf7b7