ENH: Random number generation?

[NeilGirdhar]

I'd love to have my Efax library entirely in the Array API. Would be really cool having it work with PyTorch, e.g. The only things missing now are:

• random number generation,
• JIT,
• automatic gradients, and
• Jax's "Pytree" system.

Random number generation is probably the biggest blocker though. I think that the SciPy people intended to eventually add some kind of generic random number generation. Is this the sort of thing that might find its way into this project? Seems like SciPy is the wrong home for it since they don't even intend to populate SciPy with all of the distributions they would need.

[lucascolley]

@mdhaber I wonder whether you have any thoughts on this

[lucascolley]

check out data-apis/array-api#431 if you haven't already, btw

[NeilGirdhar]

check out data-apis/array-api#431 if you haven't already, btw

I have seen it. Regarding the issues with Jax, I'm not sure they're much of a blocker. Consider an interface like this:

class RngStream:
    def __init__(self, key: KeyArray, count: JaxIntegralArray | None = None) -> None:
        super().__init__()
        if count is None:
            count = jnp.zeros((), dtype=jnp.uint32)
        else:
            assert count.ndim == 0
        self._key = key
        self._count = count

    def key(self) -> KeyArray:
        key = jr.fold_in(self._key, self._count)
        self._count += 1
        return key

    def fork(self, samples: int) -> KeyArray:
        return jr.split(self.key(), samples)

def uniform(stream: RngStream, ...) -> ...:
  ...  # etc.

This should be good enough for writing generic algorithms, and works fine for Jax provided you only pass KeyArray objects to jitted functions, which is not the responsibility of xpx.

I'm not sure, but I think RngStream above is analagous to numpy.random.SeedSequence, and KeyArray is analogous to numpy.random.Generator. Or perhaps RngStream is just Generator, and KeyArray is a Jax quirk.

[mdhaber]

There only thoughts I have are "yes, please". I tried in scipy/scipy#20549 and there was immediate opposition, but maybe someone more familiar with the JAX rng model can get around it - please see the discussion there.

[crusaderky]

From my understanding writing a generic design that includes JAX is problematic because you can't pass a mutable object - like, say, a random state - into jax.jit.

So it's impossible to write:

@jax.jit
def f(state):
    return state.uniform(123)

state = xpx.default_rng(xp=jax.numpy)  # Mocks numpy API
x = f(state) + f(state)

The above example will behave as expected for all other backends, whereas in JAX it will be the same as f(state) * 2.
What makes it worse is that there is no way for the wrapper to notice, and it's not easy to detect the loss of entropy in unit tests either.

The JAX pattern is for every function to return the new state together with the output, but that would be something that all other backends need to get on with:

@jax.jit
def f(state):
    return state.uniform(123)  # returns tuple (series, new state)

state = xpx.default_rng(xp=jax.numpy)  # state is immutable on JAX and mutable on numpy
x1, state = f(state)
x2, state = f(state)
x = x1 + x2

The alternative to the above is to have a xpx.default_rng object that can be exclusively created within a jax.jit block, cannot get out of it, and cannot get inside a nested function that is itself wrapped by jax.jit - but I question how useful that would be.

XREF https://docs.jax.dev/en/latest/jax.random.html

[NeilGirdhar]

From my understanding writing a generic design that includes JAX is problematic because you can't pass a mutable object - like, say, a random state - into jax.jit.

I addressed this in my comment above. You don't need to pass mutable generators into jitted functions.

The JAX pattern is for every function to return the new state together with the output, but that would be something that all other backends need to get on with:

I would call that the "basic Jax pattern". Many Jax users—for example users of Flax—don't use this pattern. They have a mutable generator that they either manage or have a library manage for them.

The alternative to the above is to have a xpx.default_rng object that can be exclusively created within a jax.jit block, cannot get out of it, and cannot get inside a nested function that is itself wrapped by jax.jit - but I question how useful that would be.

I don't think that's the only alternative. You can have a mutable generator that users can pass into the jitted function by converting them to a different format. Above, that's with the key and fork methods. They can then reconstruct the generator within the jitted function. This is essentially what Flax does (hidden from the user), and what I do manually when I use Equinox.

But xpx doesn't need to worry about any of that. What xpx needs to do is be able to:

• Create a mutable RNG for a namespace (using a seed and optional device)
• Sample from that RNG
• get and set state would be useful

The user can manage the rest until jitted functions are standardized.

[NeilGirdhar]

It could be something like this:

def create_generator(xp: ..., seed: int, *, device: ...) -> xpx.Generator:
  if is_jax_namespace(xp):
    return xpx.JaxGenerator(seed, device)
  ...

class Generator:  # All methods below are abstract.
  def get_state...
  def set_state...
  def uniform...
  def beta...

class JaxGenerator(Generator):  # This needs to be exposed.
  def __init__(self, seed, device):
    self._rng = jax.random.key(seed)  # move it to the device I guess

  def key(self): ...  # For passing into jitted functions.  If jitted functions are standardized, this can be made an abstract method on Generator.

  # Implement all the abstract methods.

Alternatively, you could use Generator.__new__ instead of create_generator. It would be a smaller interface, but might be a bit too magical.

[izaid]

Yes please!

[NeilGirdhar]

Here's a rough draft of how I imagine this working. This should be good enough to test a statistical library like SciPy's and might be good enough to use in a machine learning project. Comments welcome.

[crusaderky]

Here's a rough draft of how I imagine this working. This should be good enough to test a statistical library like SciPy's and might be good enough to use in a machine learning project. Comments welcome.

Sorry I must be dense - could you show a mock-up of that working with a generic function f(rng_or_state, xp: Namespace) -> Array, which may or may not be wrapped withjax.jit? I can't understand. I assume your POC is missing the pytree flattening/unflattening bit?

[NeilGirdhar]

Here's a rough draft of how I imagine this working. This should be good enough to test a statistical library like SciPy's and might be good enough to use in a machine learning project. Comments welcome.

Sorry I must be dense - could you show a mock-up of that working with a generic function f(rng_or_state, xp: Namespace) -> Array, which may or may not be wrapped withjax.jit? I can't understand. I assume your POC is missing the pytree flattening/unflattening bit?

As I said:

This should be good enough for writing generic algorithms, and works fine for Jax provided you only pass KeyArray objects to jitted functions, which is not the responsibility of xpx.

Therefore, it would be something like this:

@jax.jit
def train_one_step(model: Model, key: KeyArray) -> Model:  #  KeyArray = jax.Array
    generator = JaxGenerator(key)
    ...

def train(model: Model, generator: JaxGenerator, steps: int) -> Model:
    for i in range(steps):
        train_one_step(model, generator.key())

It is also not that hard to make a version of jax.jit that would be transparent to generators so that the above code would just be:

@xpx.jit
def train_one_step(model: Model, x: Generator, *, xp: Module) -> Model:  # Need to pass xp to help xpx.jit.
    ...

def train(model: Model, generator: Generator, steps: int) -> Model:
    for i in range(steps):
        train_one_step(model, generator, xp=jax.numpy)  # Just pass the generator!

where xpx.jit(f) means:

• Use jax.tree.map to transform every JaxGenerator into a special WrappedKeyArray.
• Apply jax.jit to f and call it.
• Usejax.tree.map to transform every WrappedKeyArray back into a JaxGenerator.

We can start another issue for generalizing the jit decorator in this way if you like? If we go down this route, we'll also need to generalize the other traced functions: vjp, jvp, grad, vmap, etc.

---

For now, I suggest we just start with an ordinary Generator and let the user sort out the traced function interfaces. What do you think?