Methods Efficiently Finding Unique Rows in a 2D Integer Array Without Using jnp.unique Along Axis?

[AndrewLiu0725]

Hi everyone,

I have a subroutine in my application that needs to be jittable and involves finding unique rows in a 2D integer array A. I’m currently running into performance issues with the jnp.unique function in JAX when using the axis argument, and I’m hoping to get some advice on a more efficient approach. Here are the properties of A:

1. A has a fixed shape (n, m), where n (the number of rows) is in the range of $10^7$ , and m (the length of each row) is around 100.
2. The elements of A are integers and lie within the range [0,k], where k is fixed.
3. Each row in A is already sorted, and no elements are repeated within a row.

What I’ve Tried So Far:
I initially tried:
unique_A, unique_index = jit(jnp.unique(A, axis=0, return_index=True, size=n, fill_value=jnp.zeros(m, dtype=jnp.int32)))

However, I found that using jnp.unique along the axis is very slow (issue #17370).

Workaround Considered:
One workaround I considered is to use a hashing function to map each row into an integer and then perform jnp.unique on the resulting 1D array. However, if the hashing function is collision-free (e.g., using powers of k+1), the resulting hash values can easily exceed the range of uint32. Since the application needs to run in single precision, partitioning the hash values into multiple uint32 segments would still require using jnp.unique along the axis, which brings us back to the original problem.

Question:
Is there an efficient, jittable algorithm that can find unique integer rows in A and avoid the performance issues with jnp.unique along the axis? Any suggestions or alternative approaches would be greatly appreciated!

Thanks in advance for your help!

[jakevdp]

First of all, there is no efficient operation in XLA to find unique values period. We have a passable implementation of np.unique based on sorting and finding adjacent equal values, but for reasons mentioned in #17370 this becomes slow for large numbers of rows because XLA's sort becomes slow for large numbers of rows.

Is there an efficient, jittable algorithm that can find unique integer rows in A and avoid the performance issues with jnp.unique along the axis? Any suggestions or alternative approaches would be greatly appreciated!

Depending on the size of your problem, you might consider the spectral analysis approach demonstrated in #17370 (comment); it's only about a dozen lines of code, it's JIT-compatible, and it should be pretty efficient. What do you think?

[AndrewLiu0725]

I'll test this spectral analysis approach and let you know if it works. However, I'm wondering about the efficiency—doesn't the complexity of jnp.linalg.eigh(L) remain $O(n^3)$? Given that, wouldn't this method still be quite expensive for large n (in my case, <10^7)?

[jakevdp]

Ah, yes if you are in a large N situation, spectral analysis will not scale well.

Perhaps you could work around this by using pure_callback to call out to np.unique?