jax.random.fold_in function is using static_argnum=(1,)

[pzhokhov]

Currently, the implementation of jax.random.fold_in function treats second argument data as static:
https://github.com/google/jax/blob/4c34541c00c02fa750b63a6ea9149909e6c4078f/jax/random.py#L193
Is that the intended behavior?
If I understand correctly, it means, for instance, that innocent-looking loops of sort:

base_key = jax.random.PRNGKey(0)
for i in range(N): 
   do_something_random(jax.random.fold_in(base_key, i))

cause fold_in to recompile at every iteration of the loop and take a lot of time.
On the other hand, if we jit-compile fold_in function, things seem to work correctly and fast; like in this gist (thanks @christopherhesse):
https://gist.github.com/christopherhesse/f493e516b7786533d76c3ef689cb6a45
Should that be the default behavior?

[mattjj]

Thanks for spotting this; I think it's almost just a bug, but it looks like we only handle 64bit Python integer data correctly by using static_argnums=1. Yet there's no test for that, and moreover we could just define fold_in only to handle 32bit integers (updating the docstring).

Want to make a PR with your fix?

[pzhokhov]

Thanks for quick response @mattjj! Not sure I understand the 64bit vs 32bit integer part... You mean 64bit integer cannot be passed as an argument to jit-compiled function, hence it was made static? I tried running the code as in the gist above, but passing np.int64(i) as a data argument; and both versions (with and without static argument) seem to work correctly.
As for making a PR - sure thing, will do!

[mattjj]

The 64bit issue is just that in the jax.random.PRNGKey function we handle 64bit values by checking isinstance and that kind of code won't work when data is traced as an argument to a jit function, just because when JAX_ENABLE_X64=0 (the default) we'll cast it down to a 32bit int. (Using static_argnums avoids that because it delays when data is effectively staged into JAX until after that isinstance check.)

I'm not worried about that behavior, so it's better to merge #1039. Thanks!

[mattjj]

I verified that this issue is solved. Thanks, @joaogui1 !