Generate non-repeative integers from a dynamic range

[zgbkdlm]

I would like to generate a fixed-size array of random uniform integers. The integers must be unique. A standard implementation is this:

def generate(n, key):
    return jax.random.choice(key, n, s, replace=False)

This funciton generate will generate an array of size s, and each element in it is a random integer and is unique.

Now I would like to jit the argument n. However, this is not possible because jax.random.choice requires that n must be static. Is there any alternative implementation/tweak?

[jakevdp]

Hi - thanks for the question! I don't know of any good answer to this. I'm not just saying that there's no function in JAX that will give you this, I'm saying I'm not sure it's possible to write an efficient implementation in JAX for what you have in mind. The current choice algorithm with replace=False pre-supposes that the number of categories is static (using either shuffle, or a gumbel top-k trick).

To select without replacement from a dynamic number of categories, I suspect you'll have to do some sort of iterative approach; something like:

vals = []
for i in range(s):
  while True:
    new_int = randint(0, n)
    if new_int not in vals:
      vals.append(new_int)
      break
return jnp.array(vals)

That's obviously not JIT compatible, but you could express the same thing using fori_loop and while_loop with a bit of thought.

I'm sure some sort of batched rejection sampling approach could improve this marginally, but the result is going to be very slow, particularly on accelerators where this kind of sequential processing is poorly supported.

[pschuh]

I think this is more jit compatible:

N = n
L = s

a = numpy.random.randint(0, N, L)
while True:
  a.sort()
  a = numpy.unique(a)
  if len(a) == L:
    break
  a = numpy.concatenate((a, numpy.random.randint(0, N, L - len(a))))

[jakevdp]

unique is not jit-compatible, unfortunately...

[zgbkdlm]

Thank you all for your answers! Ja, it is indeed hard to do it in JAX. To generate these random numbers, essentially we need to do a random permutation of jnp.arange(n) which depends on n. Rejection-like samplings apply, but the rejection ratio can be quite high depending on n. Perhaps I will for now leave the question open to see if there is any other though!

[soraros]

Maybe you could say more about how are you going to use such a dynamically shaped array? I think, inside jit, even if one can get such an array, some masking trick is still need to use it.

[jakevdp]

I think here the array is statically shaped, but the maximum value within the array is dynamic, and we don't know how to generate such an array efficiently.

[soraros]

Ah thanks! I misread the OP.

[pschuh]

Jit friendly:

a = numpy.random.randint(0, N, L)
while True:
  a.sort()
  mask = a[1:] == a[:-1]
  if mask.sum() == 0:
    break
  a = numpy.concatenate((a[:1], np.where(mask,
             numpy.random.randint(0, N, L - 1),
                                a[1:])))