Fix random seed handling in sample_hypersphere (#2688)

Summary:
Pull Request resolved: #2688

Fixes #2685

The seed was being ignored when `d=1`.

Reviewed By: esantorella

Differential Revision: D68443701

fbshipit-source-id: 2c8a5828170c8b955a452a7446d775c171db1baf

Author: saitcakmak

botorch/utils/sampling.py

@@ -163,16 +163,15 @@ def sample_hypersphere(
         >>> sample_hypersphere(d=5, n=10)
     """
     if d == 1:
-        rnd = torch.randint(0, 2, (n, 1), device=device, dtype=dtype)
+        with manual_seed(seed=seed):
+            rnd = torch.randint(0, 2, (n, 1), device=device, dtype=dtype)
         return 2 * rnd - 1
     if qmc:
         rnd = draw_sobol_normal_samples(d=d, n=n, device=device, dtype=dtype, seed=seed)
     else:
         with manual_seed(seed=seed):
-            rnd = torch.randn(n, d, dtype=dtype)
+            rnd = torch.randn(n, d, device=device, dtype=dtype)
     samples = rnd / torch.linalg.norm(rnd, dim=-1, keepdim=True)
-    if device is not None:
-        samples = samples.to(device)
     return samples
 
 

test/utils/test_sampling.py

@@ -93,6 +93,12 @@ def test_draw_sobol_samples(self):
             self.assertTrue(torch.all(samples <= bounds[1]))
             self.assertEqual(samples.device.type, self.device.type)
             self.assertEqual(samples.dtype, dtype)
+            if seed is not None:
+                # Check that seed reproduces the same samples.
+                samples2 = draw_sobol_samples(
+                    bounds=bounds, n=n, q=q, batch_shape=batch_shape, seed=seed
+                )
+                self.assertTrue(torch.equal(samples, samples2))
 
     def test_sample_simplex(self):
         for d, n, qmc, seed, dtype in itertools.product(
@@ -107,6 +113,12 @@ def test_sample_simplex(self):
             self.assertTrue(torch.max((samples.sum(dim=-1) - 1).abs()) < 1e-5)
             self.assertEqual(samples.device.type, self.device.type)
             self.assertEqual(samples.dtype, dtype)
+            if seed is not None:
+                # Check that seed reproduces the same samples.
+                samples2 = sample_simplex(
+                    d=d, n=n, qmc=qmc, seed=seed, device=self.device, dtype=dtype
+                )
+                self.assertTrue(torch.equal(samples, samples2))
 
     def test_sample_hypersphere(self):
         for d, n, qmc, seed, dtype in itertools.product(
@@ -119,6 +131,12 @@ def test_sample_hypersphere(self):
             self.assertTrue(torch.max((samples.pow(2).sum(dim=-1) - 1).abs()) < 1e-5)
             self.assertEqual(samples.device.type, self.device.type)
             self.assertEqual(samples.dtype, dtype)
+            if seed is not None:
+                # Check that seed reproduces the same samples.
+                samples2 = sample_hypersphere(
+                    d=d, n=n, qmc=qmc, seed=seed, device=self.device, dtype=dtype
+                )
+                self.assertTrue(torch.equal(samples, samples2))
 
     def test_batched_multinomial(self):
         num_categories = 5