Flipping the signs of some indices under jit

[dedeswim]

Hello!

I would like to flip the sign of an array at some indices which are computed dynamically. I tried some stuff, such as this:

import jax

@jax.jit
def flip(x, start, end):
    return x.at[start:end].mul(-1)

x = jax.numpy.arange(10)
flipped_x = flip(x, 0, 5)

but nothing works because, apparently, this creates a dynamically shaped array. Is it possible to do this in a jitable way?

I tried to take a look at lax.dynamic_slice/dynamic_update_slice as suggested by the exception thrown, but I am not sure how I could use this, as I would need to create a dynamically shaped array as update. In this toy example where start and end are known it's possible, but in my use case, start and end are not known in advance (as they change over time). I basically have to iteratively flip different regions of the array several times.

Thanks!

[jakevdp]

One pattern I've found useful in this situation is to use where with a mask:

import jax
import jax.numpy as jnp

@jax.jit
def flip(x, start, end):
  indices = jnp.arange(x.shape[0])
  return jnp.where((indices >= start) & (indices < end), x * -1, x)

x = jax.numpy.arange(10)
print(flip(x, 0, 5))
# [ 0 -1 -2 -3 -4  5  6  7  8  9]

[dedeswim]

That worked, thank you so much!