Cases not covered by the Random API

[ntessore]

What are the cases where the numpy.random.Generator interface is inadequate for sampling?

[matt-graham]

Hi @ntessore. The jrng wrapper is a nice idea, but a key issue that I would foresee is that it is incompatible with JAX's tracing mechanism and in particular will not work with it's transformations like jax.jit for just-in-time compilation.

For jax.jit specifically, there is an immediate issue that the jrng.JRNG object is not of a type that jax.jit supports as arguments to a jitted function by default. For example

import jax
import jrng

def normal(rng, loc, scale):
    return loc + rng.standard_normal() * scale

jitted_normal = jax.jit(normal)

rng = jrng.JRNG(20241216)

jitted_normal(rng, 1., 2.)

gives an error

TypeError: Error interpreting argument to <function normal at 0x7fd0c229cf40> as an abstract array. The problematic value is of type <class 'jrng.JRNG'> and was passed to the function at path rng.
This typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.

We could try to address this using the suggested approach of marking the rng argument as static (this would have a performance implication but let's ignore that for now):

jitted_normal_static = jax.jit(normal, static_argnames=("rng",))

This function is now callable, however it will always return the same value (or a shifted / scaled version if the loc and scale arguments are changed):

assert jitted_normal_static(rng, 1., 2.) == jitted_normal_static(rng, 1., 2.)

This is because jax.jit (and JAX's other transformations) expect the function being transformed to be pure without any side effects such as updates to the internal state of the arguments. In particular here when the jitted_normal_static function is first evaluated, the value of the key internally held by the JRNG object at that point will be traced and its value baked in to the traced computation so remaining fixed for all subsequent evaluations. It's this requirement of functions being pure which is a key part of why JAX's random number generator mechanism is designed in the way it is rather than matching the numpy.random API.

One option would be to have functions which take a random number generator state object to also always return the (updated) state, which is what the JAX docs recommend for using stateful computations, for example for the normal function above:

def normal(rng, loc=0.0, scale=1.0):
    return loc + scale * rng.standard_normal(), rng

To allow using the JRNG class within jitted functions you would probably also need to do something like register it as a custom pytree node type.

[ntessore]

Hi, thanks for having a look! Indeed, the jitting is the big problem here. Because even when registering the class as a custom pytree, the key array gets decoupled from the JRNG in the flatten operation, and the situation hence doesn't change (same outputs for every call).

On the one hand, I am not sure if this is fixable... Although I have not given up fully, and am still investigating.

On the other hand, this is not a problem if you do the jitting at the highest level, where the JRNG instance is created in the @jit block. This may be what you always have to do in an Array-API-like situation (as there is no @jit available to you in the implementation).

[ntessore]

I think I found a reasonable workaround for this. It is described in the README.