How to keep jax.random.choice untraced in a jitted function?

[gianlucadetommaso]

I would love to keep the output of jax.random.choice untraced in a jitted function. Here is a minimal use case:

from jax import jit, random
@jit
def fun(key):
    a = str(random.choice(key, 5))
    return len(a)
fun(random.PRNGKey(0))

With @jit, this returns 62, since a='Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>'.

Without @jit, this would return 1, since a will be an integer between 0 and 4. This is the output that I would like even when using @jit.

How can I keep the output of jax.random.choice untraced and obtain what I am looking for?

Thanks in advance!

[jakevdp]

There is no way to make random.choice execute in a non-traced manner within a jit-compiled function. That said, depending on what your actual application is (I suspect you're not actually concerned with computing the length of a string repr of a traced variable) it may be possible to do what you want to do using the new jax.pure_callback function, which requires JAX version 0.3.17 (not yet released as I write this) Make sure jax is up to date:

$ pip install jax>=0.3.17

... and then you can do something like this:

import jax
import numpy as np

def get_str_len(x):
  return np.array(len(str(x)), dtype='int32')

@jax.jit
def fun(key):
    a = jax.random.choice(key, 5)
    out_type = jax.ShapeDtypeStruct((), 'int32')
    return jax.pure_callback(get_str_len, out_type, a)
print(fun(jax.random.PRNGKey(0)))
# 1

[gianlucadetommaso]

Thanks a lot, this sounds like it should work for my application. Any indication when JAX version 0.3.17 might be released?

[jakevdp]

I just pushed the release of JAX version 0.3.17 – let me know if it works for you!

[gianlucadetommaso]

Thanks a lot! However I don't find it among the CUDA releases in here:
https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

Would it be possible to add it?

[jakevdp]

Cuda wheels are only relevant to jaxlib (the C++ support package for jax). The current release of jaxlib is version 0.3.15, which can be used with jax version 0.3.17.

[gianlucadetommaso]

I tried pure_callback and it works, but it turns out I need something a bit more complex than the simple case above. Any idea what to do when the output of the function in pure_callback (in this case get_str_len) is a generic object rather than a string?

[sharadmv]

The object would have to be registered as a Pytree. Is that not possible?

[gianlucadetommaso]

The object is a flax.core.TrainState. What should be the output type instead of jax.ShapeDtypeStruct((), 'int32') in this case?

For more clarity, I'm trying to restore a checkpoint from a random path. The path has the following form: base_path + str(i), where i is sampled from random.choice. The restoring function is jitted because called inside of lax.fori_loop. For this reason, instead of getting paths like base_path/2, I get paths like base_path/Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)>.

I would solve the problem if I sampled i from numpy.random.choice instead. But my interface provides a PRNGKeyArray. It would be very ugly not to use it. An alternative solution would be to transform the PRNGKeyArray into a numpy random seed, but I am not sure if that's possible (?).

I hope this helps clarifying rather than going too much in detail.

[jakevdp]

Since TrainState is a pytree, in place of ShapeDtypeStruct, you should pass a flax.core.TrainState object whose components have shape and dtype attributes corresponding to the expected output shapes and dtypes. This could be an actual TrainState object (if you have access to one at the call site), or you could construct a TrainState that contains ShapeDtypeStruct objects representing its contents.

[gianlucadetommaso]

It worked, thanks!

[gianlucadetommaso]

Sorry to bring this up again, but I noticed that on GPU this does not work as expected. Although the device should make no difference, the very same piece of code involving the jax.pure_callback produces different results on CPU and on GPU. I'm struggling to debug this since everything is Traced.

In principle, would there be any reason for jax.pure_callback to behave differently on CPU and GPU?

[jakevdp]

I don't know offhand of any reason that pure_callback would produce different results on GPU vs CPU. Do you have a minimal reproduction of code that shows this?

[gianlucadetommaso]

I found out that pure_callback badly interacts with flax.linen.Conv in a jit'ed function on GPU. I opened a bug issue here.

I suspect there is a bug on the Flax side of things. Maybe you have an intuition where the problem comes from?

Thank you in advance!