How do RNGs work in torchax

I’ve read the code, and it seems that we just let PyTorch manage the RNGs (I’ve tested it, and it works on eager mode), but I have a question.

I’m not entirely sure how this works, but in JAX we sometimes provide RNG keys (e.g., jax.random.key) as arguments to the loss function. Since these keys are JAX arrays, they’re tracer objects. When we generate random numbers (random boolean, for dropout mask)based on these keys, the results also become tracer objects, which are then captured in the computation graph. Shouldn’t there be an equivalent mechanism for this in the torchax?