Randomly select an element from the array that meets the given condition and return the index of the element.

[ijku]

Hello,

I'm working on a function and one of the steps is to randomly select an element from an array that meets a certain condition and return the index of the element. If the function is not decorated with jax.jit, it's easy:

import jax
import jax.numpy as jnp

def get_eligible_cells(arr, step):
    key = jax.random.PRNGKey(step)
    mask = (arr>2)
    matched_inds = np.where(jnp.reshape(mask, (-1)))[0]  # Reshape to a 1D arr and get all elements indices meet mask
    result =  jax.random.choice(key, matched_inds)
    return result

array_1d = jnp.array(([1, 2, 3], [4, 5, 6]))
ind = get_eligible_cells(array_1d, 0)
print(ind)

But if the function is decorated with jax.jit The np.where function uesed here will cause TracerArrayConversionError. I checked the JAX documentation and discovered the flaw in the above approach.

Is there an alternative?

[jakevdp]

You can use the size argument to jnp.where to keep the size of the arrays static, and then use randint to index into the valid elements:

import jax
import jax.numpy as jnp

def get_eligible_cells(arr, step):
    key = jax.random.PRNGKey(step)
    mask = (arr>2)
    matched_inds = jnp.where(jnp.reshape(mask, (-1)), size=arr.size)[0]
    idx = jax.random.randint(key, shape=(), minval=0, maxval=mask.sum())
    return matched_inds[idx]

array_1d = jnp.array(([1, 2, 3], [4, 5, 6]))
ind = jax.jit(get_eligible_cells)(array_1d, 0)
print(ind)
# 4

[ijku]

I really appreciate your detailed answer. It is exactly what I was looking for. You are amazing!