Generating random numbers in pallas on gpu

[lengstrom]

Hi, is it possible to generate random numbers in pallas on gpu? For example, triton exposes tl.language.randint to generate random integers: https://triton-lang.org/main/python-api/generated/triton.language.randint.html

[hanzhi713]

@lengstrom Curious, why did you close the issue since the feature is still not supported in the latest jax?