diff --git a/RELEASES.md b/RELEASES.md index d636e8cf8..1fb8a982f 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -11,6 +11,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) - Build the CUDA generator and the CUDA entries of `TorchBackend.__type_list__` lazily, so that using POT with CPU-only torch tensors no longer initialises a CUDA context and claims device memory (PR #847, Issue #612) - Fix the sign issue in updates of the previous transport plan in `ot.batch.proximal_bregman_log_plan_batch` (Issue #842) diff --git a/ot/backend.py b/ot/backend.py index fc087495c..f55938a93 100644 --- a/ot/backend.py +++ b/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): diff --git a/ot/sliced/_utils.py b/ot/sliced/_utils.py index e5f926fda..760214a74 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 c88ee5052..5a4570578 100644 --- a/test/test_backend.py +++ b/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)