ENH: Random number generation?

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.

@mdhaber

@mdhaber I wonder whether you have any thoughts on this

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

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.

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.

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

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.

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.

Yes please!

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.

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?

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?

In my experience, a stateful RNG is better when using Jax than sticking to Jax's bare bones approach. Therefore, I think it makes sense to go that route and tackle the jitted function problem.

I think we came pretty much to the same conclusion. For what it's worth, our proof of concept is at https://github.com/glass-dev/rng-jax. We didn't feel like we could take on and maintain a proper library at this point, so we have integrated the same thing as private API into our project (glass-dev/glass#610).

Regarding passing to jitted functions, you can't just flatten and unflatten when entering a jitted function, you need to call something like fork or key as shown here. Otherwise, you will have key reuse.

I agree. There is Generator.spawn() to achieve this within the "random API", even. However, after much experimentation, I came to the conclusion that the best way to deal with jitted functions was to prevent the wrapper RNGs from being passed to jitted functions at all, and pass a jax key directly whenever going "low-level". Code like this seems to prevent a lot of headaches:

@singledispatch
def func_generic(rng: Generator):
    ...

@func_generic.register
def _(rng: rng_jax.Generator):
    key = rng.split()
    ...  # low-level jax implementation

We discussed this in the community meeting, and the general feeling was that this wouldn't be a good fit for array-api-extra yet, as there is uncertainty regarding the requirements and ways forward.

Would you decide if this should be implemented in this repository or on a separate repository?

It sounded like everyone in the meeting would be open to work on this happening in a separate repository. If something is reached which multiple projects are interested in adopting, it also sounded like people were open to the idea of an array-api-rng repo under the organisation at that stage. I would also be open to the possibility of something getting folded in to array-api-extra at some point down the line, I just still have some scope-creep considerations.

@lucascolley

@lucascolley Then could you please setup array-api-rng?

@ntessore

@ntessore

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.

* Use`jax.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.

To avoid introducing "wrapped" decorators like xpx.jit, first and third item need to be dropped.

To drop the first, as you suggested, Generator.fork() need to return an object natively supported by jit.
To drop the 3rd, the key needs to be converted to Generator inside the function. Instead of introducing a generic function as you suggested, I would suggest to simply ~~forbid the function to take Generator~~ avoid implementing the function to make Generator as the only acceptable type and always call xpx.generator within a function (or maybe we could add decorator to do this automatically?)

def uniform_power(shape: tuple[int, ...], xp: ArrayNamespace, key_or_seed_or_none: Any = None, dtype: Any, device: Any) -> Array:
generator = xpx.generator(xp, key_or_seed_or_none)
return generator.uniform(shape=shape, dtype=dtype, device=device) ** 2

To avoid introducing "wrapped" decorators like xpx.jit, first and third item need to be dropped.

The whole point of this library is to be able to write backend-independent tests of Array API code. You could avoid xpx.jit, but then you end up having to do the behavior yourself (behavior that is a no-op for other backends).

But I think ultimately at least Jax's jit is quite special and should be exposed because, for example, within a jitted function Array API arrays can't have a device.

I don't know how jitted functions work in PyTorch though.

The whole point of this library is to be able to write backend-independent tests of Array API code.

Keep in mind that this is not the only use-case. Parts of scipy.stats which accept NumPy generators, for example, would be interested in including this in production code, not just tests.

Keep in mind that this is not the only use-case. Parts of scipy.stats which accept NumPy generators, for example, would be interested in including this in production code, not just tests.

Yes, you're absolutely right.

We should compile a clear list of motivations so as to guide the design decisions. (Similar to Joren's issue in the typing repo "What is your problem?")

@34j

Yes. I would suggest that one of you (@34j @NeilGirdhar @ntessore) makes a new repo and invites the others as collaborators, where you can work through an issue like data-apis/array-api-typing#22. If there is rough agreement on what should be worked on and a prototype going, then it might be time to create data-apis/array-api-rng.

Done: https://github.com/NeilGirdhar/array-api-rng/ I'm on vacation for the next 10 days, but I added the collaborators.

It would be nice if someone lifted the setup from array-api-typing to that repo.

@NeilGirdhar

@NeilGirdhar Thanks for creating the repo and any progress? I'm not inclined to lift the setup on my own because I don't understand their setup.

@34j

@34j You're welcome to set up the repo however you think is best. I just thought we should steal whatever ideas from their repo since I like some of the choices they made: Ruff, clean pyproject.toml, etc.

@34j You're welcome to set up the repo however you think is best. I just thought we should steal whatever ideas from their repo since I like some of the choices they made: Ruff, clean pyproject.toml, etc.

Thanks for your welcome but I won't be able to help for a while, sorry.

going to convert to a discussion — at least for now, array-api-extra is not the right place for this to be developed. Let's keep pointing to https://github.com/NeilGirdhar/array-api-rng/issues for now until we are roughly converging on a prototype.

Navigation Menu

Repository navigation

ENH: Random number generation? #308

16 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

ENH: Random number generation? #308

Description

Activity

lucascolley commented on May 31, 2025

lucascolley commented on May 31, 2025

NeilGirdhar commented on May 31, 2025

mdhaber commented on May 31, 2025

crusaderky commented on Jun 2, 2025

NeilGirdhar commented on Jun 2, 2025

NeilGirdhar commented on Jun 2, 2025

izaid commented on Jun 4, 2025

NeilGirdhar commented on Jun 5, 2025

crusaderky commented on Jun 6, 2025

NeilGirdhar commented on Jun 6, 2025

16 remaining items

ntessore commented on Jul 8, 2025

lucascolley commented on Jul 10, 2025

34j commented on Jul 11, 2025

34j commented on Jul 11, 2025

NeilGirdhar commented on Jul 11, 2025

lucascolley commented on Jul 11, 2025

NeilGirdhar commented on Jul 11, 2025

lucascolley commented on Jul 11, 2025

NeilGirdhar commented on Jul 11, 2025

34j commented on Aug 20, 2025

NeilGirdhar commented on Aug 20, 2025

34j commented on Aug 22, 2025

lucascolley commented on Aug 12, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions