orng 0.1.0

Omni RNG: A unified, cross-backend random number generator built on the Array API standard.

Omni RNG (orng)

orng provides a thin wrapper over several Array API–compatible random number generators. It mirrors a subset of the numpy.random.Generator API:

• random
• uniform
• normal
• choice
• gamma

letting you pick the underlying backend at runtime. The following backends are currently supported:

• numpy
• torch
• cupy
• jax

Installation

orng can be install from PyPI using pip:

pip install orng

Backends are optional extras that you can install as needed:

pip install "orng[numpy]"   # NumPy RNG support
pip install "orng[torch]"   # PyTorch RNG support
pip install "orng[cupy]"    # CuPy RNG support
pip install "orng[jax]"     # JAX RNG support

You can also combine extras, e.g. pip install "orng[numpy,torch]".

Quick Start

from orng import RandomGenerator

rng = RandomGenerator(backend="numpy", seed=42)
samples = rng.normal(loc=0.0, scale=1.0, size=5)
uniform = rng.uniform(low=-1.0, high=1.0, size=(2, 2))

The backend module is imported lazily. If the requested library is missing, RandomGenerator will raise an informative ImportError that points to the matching extra.

Functional Backend API

For JAX and other functional workflows, orng also provides a pure API in orng.functional:

from orng.functional import create_functional_backend

backend = create_functional_backend("numpy")
state = backend.init_state(seed=42, generator=None)

x, state = backend.normal(state, loc=0.0, scale=1.0, size=(4,), dtype=None)
y, state = backend.uniform(state, low=-1.0, high=1.0, size=(2, 2), dtype=None)

Every sampling call takes an explicit state and returns (sample, next_state). This avoids mutable RNG objects inside compiled code.

By default this API is pure (pure=True). On stateful backends (numpy, torch, and cupy) this snapshots RNG state each call. For lower overhead on those backends, you can opt into a trusted mutable fast path with pure=False:

backend = create_functional_backend("numpy", pure=False)
state = backend.init_state(seed=42, generator=None)  # numpy.random.Generator
x, state = backend.normal(state, loc=0.0, scale=1.0, size=(4,))

The JAX functional backend is always pure and does not support pure=False.

Supported functional methods:

• random
• uniform
• normal
• choice
• gamma

JAX Compilation Example

import jax
import jax.numpy as jnp
from orng.functional import create_functional_backend

backend = create_functional_backend("jax")
state = backend.init_state(seed=0, generator=None)

@jax.jit
def step(key):
    sample, next_key = backend.normal(
        key, loc=0.0, scale=1.0, size=(8,), dtype=jnp.float32
    )
    return sample, next_key

sample, state = step(state)

Functional State Reference

The functional API follows the native conventions of each backend rather than introducing a wrapper state type.

init_state(seed=..., generator=...) accepts backend-specific generator inputs:

Backend	generator argument
numpy	numpy.random.Generator
torch	torch.Generator
cupy	cupy.random.Generator
jax	JAX PRNG key array, typically from jax.random.key(...)

If generator=None, ORNG creates a new backend-native state from seed. If seed=None, the backend chooses a fresh random seed using its usual behavior.

The state value passed into random, uniform, normal, choice, and gamma also matches the backend:

Backend	pure=True state	pure=False state
numpy	NumPy bit-generator state dict	numpy.random.Generator
torch	TorchFunctionalState	torch.Generator
cupy	CuPy bit-generator state dict	cupy.random.Generator
jax	JAX PRNG key array	not supported

For example, NumPy in pure mode snapshots and returns a bit-generator state dictionary each call, while pure=False threads a numpy.random.Generator through the same functional interface. JAX always uses and returns a PRNG key.

Backend State Reference

When you pass the optional generator argument to RandomGenerator, the expected object depends on the backend:

Backend	Generator argument
numpy	numpy.random.Generator
torch	torch.Generator
cupy	cupy.random.Generator
jax	jax.random.KeyArray (from jax.random.key)

This lets you wrap an existing RNG/key instead of seeding a new one.

Citing

If you find orng useful in your work, please cite the corresponding DOI.