diff --git a/RELEASES.md b/RELEASES.md index 6ead2ce34..8b9d2902e 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/backend.py b/ot/backend.py index af622734d..eb51d171c 100644 --- a/ot/backend.py +++ b/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): diff --git a/ot/sliced/_utils.py b/ot/sliced/_utils.py index cc78a5d7c..a63d6337f 100644 --- a/ot/sliced/_utils.py +++ b/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) diff --git a/test/test_backend.py b/test/test_backend.py index 4df918140..1dc851fa6 100644 --- a/test/test_backend.py +++ b/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)