Array wrangling: permute, then truncate

[femtomc]

I have a small array wrangling problem which I think should be solvable in a JIT compat way, but I'm having issues figuring out the right idiom (especially concerning permutations).

I have a runtime index array idx - let's say it contains a single index. I have a runtime array v, with shape (100, ). I want to write a JIT compat computation which filters v to all values whose index is not idx so that filtered_v has shape (99, ) -- with all the other values retained (besides the value at idx).

So idx is not known at trace time, it's a runtime value. I tried a solution using a Boolean mask -- but that causes issues because idx is not known at trace time (so e.g. I get a NonConcreteBooleanIndexError).

Is there a simple way to construct a runtime permutation, apply it to put the value at idx at the end of the array, and then truncate via compile-time known slicing (e.g. permuted_v[:-2]) for instance?

I'm stuck on constructing the permutation. I don't think I can do any stacking (that's not trace time compat). I'm wondering if some clever combination of jnp.where will get me where I need to go.

[jakevdp]

It sounds like you want a dynamic version of jnp.delete. There's nothing like that included in JAX, but you can write it using a strategy that's something like this:

import jax
import jax.numpy as jnp

@jax.jit
def dynamic_delete(v, i):
  return jnp.where(jnp.arange(len(v) - 1) < i, v[:-1], v[1:])

v = jnp.arange(10)

print(dynamic_delete(v, 3))
# [0 1 2 4 5 6 7 8 9]

[femtomc]

this is perfect, thanks Jake

[jakevdp]

#15671 adds this logic to jnp.delete, if you pass the assume_unique_indices=True flag.

[femtomc]

@jakevdp I'm not sure if this works seamlessly when applied to arrays with more than a single dim e.g. I get:

ValueError: Incompatible shapes for broadcasting: shapes=[(49,), (49, 10), (49, 10)]

when attempting to apply dynamic_delete to an array with 2 dims.

[femtomc]

Ah, I see your implementation in #15671 is smarter.