ENH: add state_dict and load_state_dict

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

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:

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:

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:

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:

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"]

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)

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)