Efficient duplicate detection in jit functions?

[ttt733]

Let's say I've got some ShapedArray holding numbers, where all the numbers are greater than zero. I'd like to replace any duplicates in those arrays with zeros. For example:

# input
[1, 2, 3, 1, 2, 4]
#output
[1, 2, 3, 0, 0, 4]

Does anyone know of a fast way to do this on the GPU? Or is it impossible to do entirely on the GPU since you have to check the values?

[jakevdp]

With a hash table, this would be an O[N] algorithm. Unfortunately, jax does not have any hash table data structure (a deficiency it inherits from XLA, due to the fact that accelerators are not optimized for such operations) so you'll probably be limited to an O[N²] solution, in which you compare each entry to every previous entry in a vectorized fashion.

Here's one approach along those lines:

import jax.numpy as jnp

def set_duplicates_to(x, val):
  x = jnp.asarray(x)
  N, = x.shape
  assert x.ndim == 1

  # set to x-1 so columns do not match x
  M = jnp.zeros((N - 1, N)).at[:].set(x - 1)
  col = jnp.arange(N)
  row = jnp.arange(1, N)[:, None]

  # note: in JAX, overflowing indices are dropped.
  M = M.at[row, row + col].set(x)

  # compare each entry of x to every previous entry.
  return x.at[(x == M).any(0)].set(0)

x = jnp.array([1, 2, 3, 1, 2, 4])
set_duplicates_to(x, 0)
# DeviceArray([1, 2, 3, 0, 0, 4], dtype=int32)

This will not work correctly if your array is floating point and has any NaN or inf values (it requires x != x-1 and x == x to hold), but you could make it work by first converting the inputs bitwise to integers of the appropriate width.

[ttt733]

Thanks! I was trying to think of ways to get it done with a sort and a select, but this should work well enough for me anyhow. My inputs are small enough that I'm mostly just concerned with making sure it doesn't incur round trips to the CPU.

[jakevdp]

Ah, great idea! I hadn't thought of doing a sorting-based solution. Here's an approach along those lines:

import jax.numpy as jnp
from jax import lax

def set_duplicates_to(x, val):
  x = jnp.asarray(x)
  N, = x.shape
  ind = jnp.arange(N)
  x, i = lax.sort_key_val(x, ind)
  x = x.at[1:].set(jnp.where(x[1:] == x[:-1], val, x[1:]))
  return x[ind.at[i].set(ind)]

x = jnp.array([1, 2, 3, 1, 2, 4])
set_duplicates_to(x, 0)
# DeviceArray([1, 2, 3, 0, 0, 4], dtype=int32)

This relies on the fact that lax sorts are stable, so the duplicate elements remain in order in the sorted array.

[ttt733]

Nice! That looks like what I was imagining. Didn't realize where was backed by a select - I'll have to start making better use of that. Thanks again!