Dynamically indexing and updating an array in a jit operation.

[homerjed]

Hi all,

(below is an MWE for my problem)

I'm trying to operate on an array cache by dynamically setting the i-th row of this array to some value.

The i-th row is found by locating the next non-zero array to fill with the value l.

This throws a tracer error

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape int32[].

when I try to find i. This makes sense, since the value of i is not known, but how can I get around this?

Other examples simply use a list for the cache, appending instead of indexing with i but I think I cannot do this - any ideas?

Cheers!

Minimum working example here:

import jax
import jax.numpy as jnp

cache = jnp.concatenate([jnp.ones((128, 2)), jnp.zeros((128, 2))])

@jax.jit
def fn(cache, l):

    i = jnp.argmax(jnp.all(cache == 0, axis=0))

    cache = cache.at[i].set(l.squeeze())

    k = jax.lax.slice(
        cache, (0, cache.shape[-1]), (i, cache.shape[-1])
    )
    return k

l = jnp.ones((1, 2))
fn(cache, l)

Whilst it is obvious to see why this occurs - I need to be able to do the computation here

[dfm]

The fundamental issue here (as you've probably worked out) is that JAX doesn't support data-dependent shapes, so you can't return a slice of cache where the shape depends on i. The error doesn't actually occur at the i = jnp.argmax(jnp.all(cache == 0, axis=0)) line, it occurs when trying to extract the slice k.

The best approach will depend on what you actually want to do with this output. Thanks for including a MWE like this (that's always very useful!), but I don't think there's enough information here to give suggestions for alternative approaches. For example, it's not clear from your description why you need to return this dynamic slice. What about returning the full updated cache instead of a slice up to i?

[homerjed]

Thanks for the reply!

Indeed - I suspected this was fundamentally the wrong way of doing it :)

The cache refers to a collection of tokens in an attention-based model, these are cached so that - at sampling time - you have access to the tokens that were sampled before in the sequence.

I think this is why I need the index i for this operation. You may be right (I wondered this before) that I could just pass the whole cache. This would be acceptable as long as the data you're sampling is not super high-dimensional!

However, I think by doing this you can't store the cached tokens from before (you actually need to track them at each index i, although they don't need to be ordered I guess - if a bunch of them are zeros - that doesn't matter, the dot product 'ignores' them).

Just to double check - I couldn't simply concatenate these as a list that gets sequentially updated right?

Cheers!