ENH: add state and set_state methods

src/orng/backends/cupy.py

@@ -88,6 +88,14 @@ def choice(
         )
         return result
 
+    def state(self) -> Any:
+        """Return the current backend-native generator state."""
+        return self._state
+
+    def set_state(self, state: Any) -> None:
+        """Replace the current backend-native generator state."""
+        self._state = self._impl.init_state(seed=None, generator=state)
+
     def state_dict(self) -> dict[str, Any]:
         generator = self._state
         return {

src/orng/backends/jax.py

@@ -88,6 +88,14 @@ def choice(
         )
         return result
 
+    def state(self) -> Any:
+        """Return the current JAX random key."""
+        return self._state
+
+    def set_state(self, state: Any) -> None:
+        """Replace the current JAX random key."""
+        self._state = self._impl.init_state(seed=None, generator=state)
+
     def state_dict(self) -> dict[str, Any]:
         return {
             "key_data": self._impl._jnp.array(

src/orng/backends/numpy.py

@@ -88,6 +88,14 @@ def choice(
         )
         return result
 
+    def state(self) -> Any:
+        """Return the current backend-native generator state."""
+        return self._state
+
+    def set_state(self, state: Any) -> None:
+        """Replace the current backend-native generator state."""
+        self._state = self._impl.init_state(seed=None, generator=state)
+
     def state_dict(self) -> dict[str, Any]:
         generator = self._state
         return {

src/orng/backends/torch.py

@@ -94,6 +94,15 @@ def choice(
         )
         return result
 
+    def state(self) -> Any:
+        """Return the current backend-native generator state."""
+        return self._state
+
+    def set_state(self, state: Any) -> None:
+        """Replace the current backend-native generator state."""
+        self._state = self._impl.init_state(seed=None, generator=state)
+        self._impl._device = getattr(state, "device", self._impl._device)
+
     def state_dict(self) -> dict[str, Any]:
         generator = self._state
         return {

src/orng/orng.py

@@ -13,12 +13,15 @@
 
 from ._utils import SizeLike
 from .backends import create_backend, infer_backend_name_from_xp
+from .functional import FunctionalState
 
 
 class RNGBackend(Protocol):
     """Protocol representing the shim each backend must implement."""
 
-    _state: Any
+    def state(self) -> FunctionalState: ...
+
+    def set_state(self, state: FunctionalState) -> None: ...
 
     def state_dict(self) -> dict[str, Any]: ...
 
@@ -119,8 +122,6 @@ def __post_init__(self) -> None:
             device=self.device,
         )
 
-    # Public API -----------------------------------------------------------------
-
     def random(
         self,
         size: SizeLike = None,
@@ -207,7 +208,19 @@ def to_functional(
             self.backend,
             pure=pure,
         )
-        return backend, self._impl._state
+        return backend, self.state()
+
+    def state(self) -> FunctionalState:
+        """Return the current backend-native random state.
+
+        This may be a mutable generator object. Use :meth:`state_dict` when a
+        detached snapshot is required.
+        """
+        return self._impl.state()
+
+    def set_state(self, state: FunctionalState) -> None:
+        """Replace the current backend-native random state."""
+        self._impl.set_state(state)
 
     def state_dict(self) -> dict[str, Any]:
         """Return a detached, versioned snapshot of the generator state."""

tests/test_backends.py

@@ -281,6 +281,20 @@ def test_random_generator_load_state_dict(rng, backend_case):
     backend_case["assert_close"](actual, expected)
 
 
+def test_random_generator_set_state(rng, backend_case):
+    backend_name = backend_case["name"]
+    source = RandomGenerator(backend_name, seed=987)
+    expected_source = RandomGenerator(backend_name, seed=987)
+    state = source.state()
+
+    rng.set_state(state)
+    assert rng.state() is state
+    actual = rng.normal(size=(4, 3))
+    expected = expected_source.normal(size=(4, 3))
+
+    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

@@ -259,7 +259,7 @@ def choice(self, state, population, *, size, replace, probabilities):
     functional_backend, state = rng.to_functional()
 
     assert functional_backend is backend
-    assert state is rng._impl._state
+    assert state is rng.state()
 
 
 def test_array_rng_to_functional_forwards_pure(monkeypatch):