Update unit tests to work with new higher-dim SobolEngine

Summary: Account for the new capabilities of the pytorch SobolEngine (up to 21201 dims) from scipy/scipy#10844

Differential Revision: D25661131

fbshipit-source-id: de53b82c1a5a8d4bcba21cd72bbd1ccb0eec0164

test/acquisition/test_utils.py

@@ -486,7 +486,9 @@ def test_prune_inferior_points(self):
                         new_callable=mock.PropertyMock,
                     )
                 )
-                mock_event_shape.return_value = torch.Size([1, 1, 1112])
+                mock_event_shape.return_value = torch.Size(
+                    [1, 1, torch.quasirandom.SobolEngine.MAXDIM + 1]
+                )
                 es.enter_context(
                     mock.patch.object(MockPosterior, "rsample", return_value=samples)
                 )

test/sampling/test_sampler.py

@@ -420,10 +420,11 @@ def test_forward_no_collapse(self):
 
     def test_unsupported_dimension(self):
         sampler = SobolQMCNormalSampler(num_samples=2)
-        mean = torch.zeros(1112)
-        cov = DiagLazyTensor(torch.ones(1112))
+        maxdim = torch.quasirandom.SobolEngine.MAXDIM + 1
+        mean = torch.zeros(maxdim)
+        cov = DiagLazyTensor(torch.ones(maxdim))
         mvn = MultivariateNormal(mean, cov)
         posterior = GPyTorchPosterior(mvn)
         with self.assertRaises(UnsupportedError) as e:
             sampler(posterior)
-            self.assertIn("Requested: 1112", str(e.exception))
+            self.assertIn(f"Requested: {maxdim}", str(e.exception))