Repository navigation
ENH: Random number generation? #308
Description
Activity
@mdhaber I wonder whether you have any thoughts on this
- changed the title
[-]Random number generation[/-][+]ENH: Random number generation?[/+]on May 31, 2025 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
KeyArrayobjects to jitted functions, which is not the responsibility ofxpx.I'm not sure, but I think
RngStreamabove is analagous tonumpy.random.SeedSequence, andKeyArrayis analogous tonumpy.random.Generator. Or perhapsRngStreamis justGenerator, andKeyArrayis 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.
Reacted by Neil GirdharFrom 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_rngobject that can be exclusively created within ajax.jitblock, cannot get out of it, and cannot get inside a nested function that is itself wrapped byjax.jit- but I question how useful that would be.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_rngobject that can be exclusively created within ajax.jitblock, cannot get out of it, and cannot get inside a nested function that is itself wrapped byjax.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
keyandforkmethods. 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 ofcreate_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
KeyArrayobjects to jitted functions, which is not the responsibility ofxpx.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.jitthat 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.mapto transform everyJaxGeneratorinto a specialWrappedKeyArray. - Apply
jax.jittofand call it. - Use
jax.tree.mapto transform everyWrappedKeyArrayback into aJaxGenerator.
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
Generatorand let the user sort out the traced function interfaces. What do you think?- Use
16 remaining items
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
forkorkeyas 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 jaxkeydirectly 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-rngrepo 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 Then could you please setup array-api-rng?
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 byjit. - 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 takeavoid implementing the function to makeGeneratorGeneratoras the only acceptable type and always callxpx.generatorwithin 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 drop the first, as you suggested,
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
jitis 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.statswhich 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.statswhich 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?")
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.
Reacted by Lucas Colley@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 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.
Reacted by 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.
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.
- locked and limited conversation to collaborators
on Aug 12, 2026
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 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.