[MRG] Support RandomState instances in NumpyBackend.seed

RELEASES.md

@@ -15,6 +15,8 @@
 
 #### 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)
 - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859)
 - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)
 - Fix device placement in `ot.batch.bregman_projection_batch` so `ot.solve_batch(..., method="sinkhorn")` no longer crashes on GPU when the torch default device is CPU (PR #851)

ot/backend.py

@@ -1456,7 +1456,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

@@ -852,6 +852,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)