JAX tracing out the random global key

[RustyBamboo]

I was using jax backend + jax jit to optimize parameters of a quimb tensor network but with some simulated noise. Noticed that the outputs were not random, and I believe it is due to jit tracing out the global key in autoray.

Minimal example:

import jax
import autoray as ar

@jax.jit
def sampler():
    return ar.do("random.normal", like='jax')
    # return jax.random.normal(key=ar.autoray.jax_random_get_key())

for _ in range(10):
    print(sampler()) 

0.11975035
0.11975035
0.11975035
0.11975035
0.11975035
0.11975035
0.11975035
0.11975035
0.11975035
0.11975035

Any ideas of a clean workaround that can keep code backend-agnostic? I guess this fundamentally breaks the assumption of a changing global RNG state that all the other backends have, but I would like to keep jit as it does massively increase the code's performance.

[jcmgray]

Yes I think the best thing to do here would be for autoray to provide a unified random.default_rng interface, a la numpy, and encourage that for all random use along with requiring an explicit seed input for jax. Since jax fundamentally requires pure functions I'm not sure there is any other way.

@ar.autojit
def foo(x, seed):
    rng = ar.do("random.default_rng", seed=seed, like=x)
    return rng.normal(size=(2, 3, 4))

I do have an implementation of this kicking around but haven't pushed it to autoray yet.

[RustyBamboo]

Sweet, I think that nicely solves this.

Huge thanks for all the amazing work you've put into these tools!

[jcmgray]

I have added an initial version of this for jax and torch if you want to try, it will take a bit more effort/time to incorporate into quimb itself however.

[RustyBamboo]

Awesome, thank you!