How to construct a 'jit'able function to delete an entry of an array

[abhiroop513]

I was trying to write the following 'jit'able function to delete some entries of an array:

import jax.numpy as jnp
from jax import jit

x = jnp.array([1., 2., 3., 4., 5., 6., 7., 8., 9.])

def zdel(x, indx):
  return jnp.delete(x, obj=indx)

zjit = jit(zdel)

indx = 5
z = zjit(x, indx)
print(z)

This gives ConcretizationType Error for the variable "indx".
(Although zdel works without jit )

How will I write a 'jit'able function for the same

[jakevdp]

The index argument to delete must be static, so the easiest way to jit-compile this function is to mark the argument as static:

zjit = jit(zdel, static_argnames=['indx'])

If you need to use a dynamic index, you can often compute the same result using operations that are compatible with dynamic values. Here's one way you might do it:

@jit
def zdel2(x, indx):
  return jnp.where(jnp.arange(len(result) - 1) < indx, x[:-1], x[1:])

print(zjit(x, indx))
# [1. 2. 3. 4. 5. 7. 8. 9.]

[fuglede]

I'm in the last boat, needing to remove entries dynamically, so this is very helpful!

Moreover, I need to do it quite a few times; say I start with an array of length 1000, then I iteratively compute an index (in some non-trivial fashion) and remove it from the array, and I perform this iteration n times (and n = 100 in my particular case). So I'm curious what happens internally: Will the JIT compiler end up generating a separate path for each of those n iterations, so that the generated code has size that grows exponentially in n? I assume not, and what comes out at the end is certainly still quite fast but it does seem to take quite a while to compile. Is there a better approach than jnp.where in this case?

[jakevdp]

Note that in recent versions of JAX, jnp.delete(x, i) is now compatible with JIT so long as i is a scalar, so you shouldn't need this workaround anymore.

As for your followup question, it sounds like you're using a for loop to call delete repeatedly on an array. Unrelated to how delete is implemented, this will generally cause very long compilations because JAX effectively unwraps Python control flow. The best way around this would probably be to restructure your computation in a way that's compatible with fori_loop or other structured control flow primitives.

A comment on a long-resolved issue is not the best place to go into more detail on this – I'd suggest opening a new discussion thread if you need more help with this.

[fuglede]

Gotcha, I created a new discussion thread to follow up.