From 9a281dc19c0ea9f0731fb60b1ffb80d892ff719b Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:22:00 +0100 Subject: [PATCH 1/3] ENH: add state_dict and load_state_dict --- src/orng/backends/cupy.py | 30 +++++++++++++++++ src/orng/backends/jax.py | 24 ++++++++++++++ src/orng/backends/numpy.py | 30 +++++++++++++++++ src/orng/backends/torch.py | 21 ++++++++++++ src/orng/orng.py | 68 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 173 insertions(+) diff --git a/src/orng/backends/cupy.py b/src/orng/backends/cupy.py index 15eeaf0..e537ea1 100644 --- a/src/orng/backends/cupy.py +++ b/src/orng/backends/cupy.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +from collections.abc import Mapping from typing import Any from .._utils import SizeLike @@ -87,6 +88,35 @@ def choice( ) return result + def state_dict(self) -> dict[str, Any]: + generator = self._state + return { + "bit_generator": type(generator.bit_generator).__name__, + "bit_generator_state": copy.deepcopy( + generator.bit_generator.state + ), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + try: + bit_generator_name = state["bit_generator"] + bit_generator_state = state["bit_generator_state"] + except KeyError as exc: + raise ValueError( + "CuPy RNG state requires 'bit_generator' and " + "'bit_generator_state'." + ) from exc + bit_generator_cls = getattr( + self._impl._cupy.random, bit_generator_name, None + ) + if bit_generator_cls is None: + raise ValueError( + f"Unknown CuPy bit generator '{bit_generator_name}'." + ) + bit_generator = bit_generator_cls() + bit_generator.state = copy.deepcopy(bit_generator_state) + self._state = self._impl._cupy.random.Generator(bit_generator) + class CuPyFunctionalBackend: def __init__(self, *, pure: bool = True) -> None: diff --git a/src/orng/backends/jax.py b/src/orng/backends/jax.py index 792cf8c..e5eb93b 100644 --- a/src/orng/backends/jax.py +++ b/src/orng/backends/jax.py @@ -1,6 +1,7 @@ from __future__ import annotations import secrets +from collections.abc import Mapping from typing import Any from .._utils import SizeLike, normalize_shape @@ -87,6 +88,29 @@ def choice( ) return result + def state_dict(self) -> dict[str, Any]: + return { + "key_data": self._impl._jnp.array( + self._impl._jax.random.key_data(self._state) + ), + "implementation": str( + self._impl._jax.random.key_impl(self._state) + ), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + try: + key_data = state["key_data"] + implementation = state["implementation"] + except KeyError as exc: + raise ValueError( + "JAX RNG state requires 'key_data' and 'implementation'." + ) from exc + self._state = self._impl._jax.random.wrap_key_data( + self._impl._jnp.array(key_data), + impl=implementation, + ) + class JAXFunctionalBackend: def __init__(self) -> None: diff --git a/src/orng/backends/numpy.py b/src/orng/backends/numpy.py index 3914d51..34e61a2 100644 --- a/src/orng/backends/numpy.py +++ b/src/orng/backends/numpy.py @@ -1,6 +1,7 @@ from __future__ import annotations import copy +from collections.abc import Mapping from typing import Any from .._utils import SizeLike @@ -87,6 +88,35 @@ def choice( ) return result + def state_dict(self) -> dict[str, Any]: + generator = self._state + return { + "bit_generator": type(generator.bit_generator).__name__, + "bit_generator_state": copy.deepcopy( + generator.bit_generator.state + ), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + try: + bit_generator_name = state["bit_generator"] + bit_generator_state = state["bit_generator_state"] + except KeyError as exc: + raise ValueError( + "NumPy RNG state requires 'bit_generator' and " + "'bit_generator_state'." + ) from exc + bit_generator_cls = getattr( + self._impl._np.random, bit_generator_name, None + ) + if bit_generator_cls is None: + raise ValueError( + f"Unknown NumPy bit generator '{bit_generator_name}'." + ) + bit_generator = bit_generator_cls() + bit_generator.state = copy.deepcopy(bit_generator_state) + self._state = self._impl._np.random.Generator(bit_generator) + class NumPyFunctionalBackend: def __init__(self, *, pure: bool = True) -> None: diff --git a/src/orng/backends/torch.py b/src/orng/backends/torch.py index b219a58..771b7e8 100644 --- a/src/orng/backends/torch.py +++ b/src/orng/backends/torch.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import Any @@ -93,6 +94,26 @@ def choice( ) return result + def state_dict(self) -> dict[str, Any]: + generator = self._state + return { + "generator_state": generator.get_state().clone(), + "device": str(getattr(generator, "device", self._impl._device)), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + try: + generator_state = state["generator_state"] + except KeyError as exc: + raise ValueError( + "Torch RNG state requires 'generator_state'." + ) from exc + device = state.get("device") or "cpu" + self._impl._device = self._impl._torch.device(device) + generator = self._impl._torch.Generator(device=device) + generator.set_state(generator_state.clone()) + self._state = generator + @dataclass(frozen=True) class TorchFunctionalState: diff --git a/src/orng/orng.py b/src/orng/orng.py index 3a0ed36..6d422ad 100644 --- a/src/orng/orng.py +++ b/src/orng/orng.py @@ -7,6 +7,7 @@ from __future__ import annotations +from collections.abc import Mapping from dataclasses import dataclass from typing import Any, Protocol @@ -19,6 +20,10 @@ class RNGBackend(Protocol): _state: Any + def state_dict(self) -> dict[str, Any]: ... + + def load_state_dict(self, state: Mapping[str, Any]) -> None: ... + def random(self, *, size: SizeLike, dtype: Any | None) -> Any: ... def uniform( @@ -204,6 +209,41 @@ def to_functional( ) return backend, self._impl._state + def state_dict(self) -> dict[str, Any]: + """Return a detached, versioned snapshot of the generator state.""" + backend = self.backend.lower() + if backend == "pytorch": + backend = "torch" + return { + "version": 1, + "backend": backend, + "state": self._impl.state_dict(), + } + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + """Restore this generator from :meth:`state_dict` output.""" + _, backend, backend_state = _validate_state_dict(state) + current_backend = self.backend.lower() + if current_backend == "pytorch": + current_backend = "torch" + if backend != current_backend: + raise ValueError( + f"Cannot load '{backend}' RNG state into a " + f"'{current_backend}' generator." + ) + self._impl.load_state_dict(backend_state) + if backend == "torch": + self.device = backend_state.get("device") + + @classmethod + def from_state_dict(cls, state: Mapping[str, Any]) -> "RandomGenerator": + """Construct a generator from :meth:`state_dict` output.""" + _, backend, backend_state = _validate_state_dict(state) + device = backend_state.get("device") if backend == "torch" else None + rng = cls(backend=backend, device=device) + rng._impl.load_state_dict(backend_state) + return rng + class ArrayRNG(RandomGenerator): """Deprecated alias for :class:`RandomGenerator`.""" @@ -219,4 +259,32 @@ def __post_init__(self) -> None: super().__post_init__() +def _validate_state_dict( + state: Mapping[str, Any], +) -> tuple[int, str, Mapping[str, Any]]: + if not isinstance(state, Mapping): + raise TypeError("ORNG state must be a mapping.") + missing = { + key for key in ("version", "backend", "state") if key not in state + } + if missing: + raise ValueError( + "ORNG state is missing required keys: " + + ", ".join(sorted(missing)) + ) + version = state["version"] + if version != 1: + raise ValueError(f"Unsupported ORNG state version {version!r}.") + backend = state["backend"] + if not isinstance(backend, str): + raise TypeError("ORNG state 'backend' must be a string.") + backend = backend.lower() + if backend == "pytorch": + backend = "torch" + backend_state = state["state"] + if not isinstance(backend_state, Mapping): + raise TypeError("ORNG backend state must be a mapping.") + return version, backend, backend_state + + __all__ = ["ArrayRNG", "RandomGenerator"] From 8ed109b86fc9f3d50c9bc64028b34bf7771f8d7b Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:22:28 +0100 Subject: [PATCH 2/3] TST: add tests for state dicts --- tests/test_backends.py | 23 ++++++++++++++++++ tests/test_orng.py | 54 ++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 75 insertions(+), 2 deletions(-) diff --git a/tests/test_backends.py b/tests/test_backends.py index 6a7282b..a8681a5 100644 --- a/tests/test_backends.py +++ b/tests/test_backends.py @@ -258,6 +258,29 @@ def test_array_rng_normal(rng): assert x.shape == (4, 4) +def test_random_generator_state_round_trip(rng, backend_case): + rng.normal(size=(3, 2)) + state = rng.state_dict() + + expected = rng.normal(size=(4, 3)) + restored = RandomGenerator.from_state_dict(state) + actual = restored.normal(size=(4, 3)) + + backend_case["assert_close"](actual, expected) + + +def test_random_generator_load_state_dict(rng, backend_case): + rng.uniform(size=(5,)) + state = rng.state_dict() + expected = rng.uniform(size=(5,)) + + rng.uniform(size=(5,)) + rng.load_state_dict(state) + actual = rng.uniform(size=(5,)) + + backend_case["assert_close"](actual, expected) + + def test_array_rng_gamma(rng): x = rng.gamma(shape=2.0, scale=2.0, size=(6, 2)) assert x.shape == (6, 2) diff --git a/tests/test_orng.py b/tests/test_orng.py index 516ffa4..517124d 100644 --- a/tests/test_orng.py +++ b/tests/test_orng.py @@ -286,8 +286,10 @@ def choice(self, state, population, *, size, replace, probabilities): monkeypatch.setattr( "orng.functional.create_functional_backend", - lambda name, pure: captured.update({"name": name, "pure": pure}) - or DummyFunctionalBackend(), + lambda name, pure: ( + captured.update({"name": name, "pure": pure}) + or DummyFunctionalBackend() + ), ) rng = RandomGenerator(backend="numpy", seed=123) @@ -296,6 +298,54 @@ def choice(self, state, population, *, size, replace, probabilities): assert captured == {"name": "numpy", "pure": False} +def test_random_generator_state_dict_is_detached(): + np = pytest.importorskip("numpy") + rng = RandomGenerator(backend="numpy", seed=123) + rng.normal(size=3) + + state = rng.state_dict() + expected = RandomGenerator.from_state_dict(state).normal(size=4) + rng.normal(size=20) + actual = RandomGenerator.from_state_dict(state).normal(size=4) + + np.testing.assert_allclose(actual, expected) + + +@pytest.mark.parametrize( + ("state", "error", "match"), + [ + ({}, ValueError, "missing required keys"), + ( + {"version": 99, "backend": "numpy", "state": {}}, + ValueError, + "Unsupported ORNG state version", + ), + ( + {"version": 1, "backend": 1, "state": {}}, + TypeError, + "backend.*string", + ), + ( + {"version": 1, "backend": "numpy", "state": "bad"}, + TypeError, + "backend state.*mapping", + ), + ], +) +def test_random_generator_rejects_invalid_state(state, error, match): + with pytest.raises(error, match=match): + RandomGenerator.from_state_dict(state) + + +def test_random_generator_rejects_state_for_different_backend(): + rng = RandomGenerator(backend="numpy", seed=123) + state = rng.state_dict() + state["backend"] = "jax" + + with pytest.raises(ValueError, match="Cannot load 'jax'"): + rng.load_state_dict(state) + + def test_numpy_backend_import_error_mentions_extra(monkeypatch): # Ensure numPy import attempts go through builtins.__import__ monkeypatch.delitem(sys.modules, "numpy", raising=False) From bb5c5e6a79b521e7a62669dcb0cae1d5a714e284 Mon Sep 17 00:00:00 2001 From: "Michael J. Williams" Date: Mon, 13 Jul 2026 17:22:59 +0100 Subject: [PATCH 3/3] DOC: mention state in readme --- README.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/README.md b/README.md index 8975f5d..1dad799 100644 --- a/README.md +++ b/README.md @@ -50,6 +50,25 @@ The backend module is imported lazily. If the requested library is missing, `RandomGenerator` will raise an informative `ImportError` that points to the matching extra. +## Saving and restoring state + +`RandomGenerator` exposes a versioned state mapping for checkpointing: + +```python +state = rng.state_dict() + +restored = RandomGenerator.from_state_dict(state) +assert restored.backend == rng.backend +``` + +An existing generator can also be reset in place: + +```python +rng.load_state_dict(state) +``` + +The returned mapping is detached from the live generator. It contains native backend arrays where applicable. + ## Functional Backend API For JAX and other functional workflows, `orng` also provides a pure API in