[MRG] Support RandomState instances in NumpyBackend.seed

RELEASES.md

@@ -8,6 +8,7 @@
 
 #### Closed issues
 
+- Allow `NumpyBackend.seed` to adopt an existing `np.random.RandomState` instance and remove NumPy-specific random sampling paths in sliced utilities (PR #849, Issue #848)
 - Preserve input dtype and device for expected sliced plans, avoid materializing dense distance matrices for sparse plans, and fix weighted sparse-distance ordering (PR #846, Issue #845)
 - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (Issue #842)
 - Load triton before TensorFlow in `ot.backend` so that building a torch optimizer no longer segfaults the interpreter, and remove the `torch<2.12` pin from the doctest and documentation requirements (PR #839, Issue #816)

ot/backend.py

@@ -1430,7 +1430,9 @@ def reshape(self, a, shape):
         return np.reshape(a, shape)
 
     def seed(self, seed=None):
-        if seed is not None:
+        if isinstance(seed, np.random.RandomState):
+            self.rng_ = seed
+        elif seed is not None:
             self.rng_.seed(seed)
 
     def rand(self, *size, type_as=None):

ot/sliced/_utils.py

@@ -54,12 +54,9 @@ def get_random_projections(d, n_projections, seed=None, backend=None, type_as=No
     else:
         nx = backend
 
-    if isinstance(seed, np.random.RandomState) and str(nx) == "numpy":
-        projections = seed.randn(d, n_projections)
-    else:
-        if seed is not None:
-            nx.seed(seed)
-        projections = nx.randn(d, n_projections, type_as=type_as)
+    if seed is not None:
+        nx.seed(seed)
+    projections = nx.randn(d, n_projections, type_as=type_as)
 
     projections = projections / nx.sqrt(nx.sum(projections**2, 0, keepdims=True))
     return projections
@@ -99,12 +96,9 @@ def get_projections_sphere(d, n_projections, seed=None, backend=None, type_as=No
     else:
         nx = backend
 
-    if isinstance(seed, np.random.RandomState) and str(nx) == "numpy":
-        Z = seed.randn(n_projections, d, 2)
-    else:
-        if seed is not None:
-            nx.seed(seed)
-        Z = nx.randn(n_projections, d, 2, type_as=type_as)
+    if seed is not None:
+        nx.seed(seed)
+    Z = nx.randn(n_projections, d, 2, type_as=type_as)
 
     projections, _ = nx.qr(Z)
     return projections
@@ -159,12 +153,9 @@ def get_random_rotations(d, n_rotations, seed=None, backend=None, type_as=None):
     else:
         nx = backend
 
-    if isinstance(seed, np.random.RandomState) and str(nx) == "numpy":
-        Z = seed.randn(n_rotations, d, d)
-    else:
-        if seed is not None:
-            nx.seed(seed)
-        Z = nx.randn(n_rotations, d, d, type_as=type_as)
+    if seed is not None:
+        nx.seed(seed)
+    Z = nx.randn(n_rotations, d, d, type_as=type_as)
 
     Q, R = nx.qr(Z)
     diagonal = nx.sum(R * nx.eye(d, type_as=R)[None, :, :], axis=-1)

test/test_backend.py

@@ -840,6 +840,17 @@ def test_random_backends(nx):
         res = nx.randperm(size=[5, 12])
 
 
+def test_numpy_backend_seed_random_state():
+    nx = ot.backend.NumpyBackend()
+    rng = np.random.RandomState(42)
+    expected_rng = np.random.RandomState(42)
+
+    nx.seed(rng)
+
+    assert nx.rng_ is rng
+    np.testing.assert_array_equal(nx.randn(5, 2), expected_rng.randn(5, 2))
+
+
 def test_gradients_backends():
     rnd = np.random.RandomState(0)
     v = rnd.randn(10)