Apply function only on slice of array under jit

[fedetask]

I want to perform an operation like

def fun(x, index):
    x[:index] = other_fun(x[:index])
    return x

This cannot be performed under jit. Is there a way of doing this with jax.ops or jax.lax?
I thought of using jax.ops.index_update(x, idx, y) but I cannot find a way of computing y without incurring in the same problem again.

[lamflokas]

Fundamentally there are two different problems. The first one is that jax does not support array updates with the x{:index] syntax. This limitation applies in both jitted and non-jitted functions. As you said, jax provides an alternative based on jax.ops.index.update. This is the correct way to handle this issue.

The other problem is that jitted code does not support dynamically shaped arrays and x[:index] inside other_fun is dynamically shaped because we do not know the value of index at jitting time. This limitation is fundamental because the underlying optimizations of jax need to know the shapes of intermediates in order to speed up the evaluation. If the function is called many times with the same index value it may make sense to jit a specialized version of fun for the fixed index.

import functools
import jax
import jax.numpy as np

def fun(index, x):
   y  = other_fun(x[:index])
   x  = jax.ops.index_update(x, index, y)
   return x

# Assume that fun will be called many times 
# with index=2 in the future
q = jax.jit(functools.partial(fun, 2))
print(q(np.arange(5))

[jakevdp]

x-ref https://stackoverflow.com/q/68419632; I'll repeat the answer here for completeness.

For a dynamic index, you can do this using jnp.where along with masking. For example:

import jax
import jax.numpy as jnp

def other_fun(x):
    return x + 1

@jax.jit
def fun(x, index):
  mask = jnp.arange(x.shape[0]) < index
  return jnp.where(mask, other_fun(x), x)

x = jnp.arange(5)
print(fun(x, 3))
# [1 2 3 3 4]

[pharringtonp19]

Is it possible to extend the solution to the case where we want to return the output of other_fun evaluated only on a slice of x?

[jakevdp]

No, it is not possible to return arrays of dynamic shape from a JIT-compiled function. That said, if you mark index as static, then you could return x[:index]. for example:

from jax import jit, partial
import jax.numpy as jnp

@partial(jit, static_argnums=1)
def f(x, index):
  return x[:index] ** 2

f(jnp.arange(100), 5)
# DeviceArray([ 0,  1,  4,  9, 16], dtype=int32)

But there is no way to do this within a jit-compiled function if index is dynamic.

[pharringtonp19]

@jakevdp Thanks for the help

Is there anyway to apply vmap over the index? I know Mattj's HashableArrayWrapper allows us to use jnp.array for the static argument, but I haven't been able to make this work

[jakevdp]

I don't believe there is: vmap works via a similar tracing method to JIT, and mapped arguments cannot be static.

At a higher level, it sounds like you are interested in a ragged array data structure, which XLA/JAX does not currently support.

[pharringtonp19]

@jakevdp I think that's exactly what I was looking for.

I am trying to make a jit/vmap-able function that splits a data set into a training and test set. My aim is to train models in parallel across devices using varying fractions of the original dataset for the training set.