Functional or stateful API

@NeilGirdhar

There are two possible ways to implement a generic random number generation API:

Functional (jax-like), where the key is manually advanced by the user, and random number generation is stateless;
Stateful (numpy-like), where the random state is kept in a Generator object that can generate random numbers.

I played out both approaches a while ago here: data-apis/array-api#874

The summary of that discussion was that

the functional API can be adopted for numpy with full generality,
the stateful API can be adopted for jax with some caveats (i.e. no state changes on device).

There are toy models of both available:

Functional API for numpy: https://github.com/ntessore/numpy_random_api
Generator API for jax: https://github.com/glass-dev/rng-jax

For what it's worth, we are using the latter approach in production here: https://github.com/glass-dev/glass/blob/main/glass/jax.py
But that was done because we were already using numpy Generators, not because the approach is better (which I don't believe).

NB: @NeilGirdhar, my invitation to this repo was expired, but if you could invite me again then I will move the above discussion here.

NB: @NeilGirdhar, my invitation to this repo was expired, but if you could invite me again then I will move the above discussion here.

Done!

I also went with the stateful approach, which I adapted from flax.

Feel free to let me know if you need any additional permissions 😄

Thanks Neil! I don't think I need anything :)

On topic: The stateful API is surely more attractive to a wider base of users who are accustomed to writing NumPy code, which is why we went for it. But there remains the fact that it cannot easily be moved to devices in jax, which I think has the potential to catch people out in the worst possible way. So there's a real decision to be made here.

But there remains the fact that it cannot easily be moved to devices in jax, which I think has the potential to catch people out in the worst possible way. So there's a real decision to be made here.

IIRC, the issue is that it can't cross jit/autograd boundaries? If I remember, the stateful class I went with is neither a pytree nor hashable, so they will get an error if they try to pass it across the boundary. You would have to really try to get it across the boundary.

I agree that the stateful approach is more intuitive to most people. But also, it's much easier to avoid the common error of key reuse. Also, it lets you write simpler code. (No premature splitting.)

These three advantages coupled with the fact that the main downside is hard to actually have happen to you makes the stateful approach more convincing to me. WDYT?

IIRC, the issue is that it can't cross jit/autograd boundaries? If I remember, the stateful class I went with is neither a pytree nor hashable, so they will get an error if they try to pass it across the boundary. You would have to really try to get it across the boundary.

That's right and I went down the same path. That being said, it is a rare but real pain point when you have to reorganise code around the constraint that "random number generation must be high level".

around the constraint that "random number generation must be high level".

I don't like reorganizing code in unintuitive ways, so I went with this pattern:

@jit
def g(keys: Mapping[str, KeyArray]):
  streams = create_streams(keys)
  # Now use streams in the usual stateful way:
  jax.random.uniform(streams["inference"].key())
  jax.random.normal(streams["inference"].key())
  # etc.

def f(streams: Mapping[str, RngStream]):
  keys = sample_streams(streams)
  g(keys)

might it be nice to draft adding one of your approaches to this repo, so we can compare it with alternative approaches inline?

Navigation Menu

Repository navigation

Functional or stateful API #4

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Functional or stateful API #4

Description

Activity

NeilGirdhar commented on Jun 7, 2026

NeilGirdhar commented on Jun 7, 2026

NeilGirdhar commented on Jun 7, 2026

ntessore commented on Jun 7, 2026

NeilGirdhar commented on Jun 7, 2026

ntessore commented on Jun 8, 2026

NeilGirdhar commented on Jun 8, 2026

lucascolley commented on Aug 12, 2026

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions