Fix tensor shapes in unit tests (#1620)

Summary:
Pull Request resolved: #1620

As a result of switching from `self.AssertTrue(torch.allclose(...))` to `self.AssertAllClose` in unit tests, we will now also have checks that tensors compared are the same shape and not just numerically equal. Some of our current tests were failing; this fixes that by changing the shapes of the compared tensors.

Reviewed By: Balandat

Differential Revision: D42402387

fbshipit-source-id: 2a375b21d1c309142cfd40a231a5f5bf4861be5f

test/acquisition/test_analytic.py

@@ -143,7 +143,7 @@ def test_expected_improvement(self):
                 model=mm, best_f=0.0, posterior_transform=transform
             )
             X = torch.rand(1, 2, device=self.device, dtype=dtype)
-            ei_expected = torch.tensor(0.6910, device=self.device, dtype=dtype)
+            ei_expected = torch.tensor([0.6910], device=self.device, dtype=dtype)
             self.assertTrue(torch.allclose(ei(X), ei_expected, atol=1e-4))
             self.assertTrue(torch.allclose(log_ei(X), ei_expected.log(), atol=1e-4))
 
@@ -374,13 +374,13 @@ def test_upper_confidence_bound(self):
             module = UpperConfidenceBound(model=mm, beta=1.0)
             X = torch.zeros(1, 1, device=self.device, dtype=dtype)
             ucb = module(X)
-            ucb_expected = torch.tensor([1.5], device=self.device, dtype=dtype)
+            ucb_expected = torch.tensor(1.5, device=self.device, dtype=dtype)
             self.assertTrue(torch.allclose(ucb, ucb_expected, atol=1e-4))
 
             module = UpperConfidenceBound(model=mm, beta=1.0, maximize=False)
             X = torch.zeros(1, 1, device=self.device, dtype=dtype)
             ucb = module(X)
-            ucb_expected = torch.tensor([0.5], device=self.device, dtype=dtype)
+            ucb_expected = torch.tensor(0.5, device=self.device, dtype=dtype)
             self.assertTrue(torch.allclose(ucb, ucb_expected, atol=1e-4))
 
             # check for proper error if multi-output model
@@ -447,7 +447,7 @@ def test_constrained_expected_improvement(self):
             X = torch.empty(1, 1, device=self.device, dtype=dtype)  # dummy
             ei = module(X)
             ei_expected_unconstrained = torch.tensor(
-                0.19780, device=self.device, dtype=dtype
+                [0.19780], device=self.device, dtype=dtype
             )
             ei_expected = ei_expected_unconstrained * 0.5
             self.assertTrue(torch.allclose(ei, ei_expected, atol=1e-4))
@@ -508,7 +508,7 @@ def test_constrained_expected_improvement(self):
             X = torch.empty(1, 1, device=self.device, dtype=dtype)  # dummy
             ei = module(X)
             ei_expected_unconstrained = torch.tensor(
-                0.19780, device=self.device, dtype=dtype
+                [0.19780], device=self.device, dtype=dtype
             )
             ei_expected = ei_expected_unconstrained * 0.5 * 0.5 * 0.5
             self.assertTrue(torch.allclose(ei, ei_expected, atol=1e-4))
@@ -527,7 +527,7 @@ def test_constrained_expected_improvement(self):
             log_module_min = LogConstrainedExpectedImprovement(**kwargs)
             ei_min = module_min(X)
             ei_expected_unconstrained_min = torch.tensor(
-                0.6978, device=self.device, dtype=dtype
+                [0.6978], device=self.device, dtype=dtype
             )
             ei_expected_min = ei_expected_unconstrained_min * 0.5
             self.assertTrue(torch.allclose(ei_min, ei_expected_min, atol=1e-4))

test/acquisition/test_knowledge_gradient.py

@@ -204,7 +204,8 @@ def test_evaluate_q_knowledge_gradient(self):
                     patch_f.assert_called_once()
                     cargs, ckwargs = patch_f.call_args
                     self.assertEqual(ckwargs["X"].shape, torch.Size([1, 3, 1]))
-            self.assertTrue(torch.allclose(val, mean.mean() - current_value, atol=1e-4))
+            expected = (mean.mean() - current_value).reshape([])
+            self.assertTrue(torch.allclose(val, expected, atol=1e-4))
             self.assertTrue(torch.equal(qKG.extract_candidates(X), X[..., :-n_f, :]))
             # test objective (inner MC sampling)
             objective = GenericMCObjective(objective=lambda Y, X: Y.norm(dim=-1))
@@ -246,7 +247,7 @@ def test_evaluate_q_knowledge_gradient(self):
                     patch_f.assert_called_once()
                     cargs, ckwargs = patch_f.call_args
                     self.assertEqual(ckwargs["X"].shape, torch.Size([1, 1, 1]))
-                    val_expected = (mean * weights).sum(-1).mean(0)
+                    val_expected = (mean * weights).sum(-1).mean(0)[0]
                     self.assertTrue(torch.allclose(val, val_expected))
 
     def test_evaluate_kg(self):
@@ -478,7 +479,7 @@ def test_evaluate_qMFKG(self):
                     patch_f.assert_called_once()
                     cargs, ckwargs = patch_f.call_args
                     self.assertEqual(ckwargs["X"].shape, torch.Size([1, 16, 4]))
-                    val_exp = torch.tensor([1.0], device=self.device, dtype=dtype)
+                    val_exp = torch.tensor(1.0, device=self.device, dtype=dtype)
                     self.assertTrue(torch.allclose(val, val_exp, atol=1e-4))
 
     def test_fixed_evaluation_qMFKG(self):

test/models/kernels/test_categorical.py

@@ -69,7 +69,7 @@ def test_ard(self):
         kernel.eval()
 
         sc_dists = x1.unsqueeze(-2) != x2.unsqueeze(-3)
-        sc_dists = sc_dists / lengthscales.unsqueeze(-2)
+        sc_dists = sc_dists / lengthscales
         actual = torch.exp(-sc_dists.mean(-1))
         res = kernel(x1, x2).to_dense()
         self.assertTrue(torch.allclose(res, actual))

test/models/test_fully_bayesian.py

@@ -679,7 +679,7 @@ def potential_fn(z):
         grads, val = potential_grad(potential_fn, z)
         self.assertTrue(torch.allclose(grads["K"], -0.5 * torch.eye(2)))
         norm_mvn = torch.distributions.Normal(0, 1)
-        self.assertTrue(torch.allclose(val, 2 * norm_mvn.log_prob(torch.zeros(1))))
+        self.assertTrue(torch.allclose(val, 2 * norm_mvn.log_prob(torch.tensor(0.0))))
 
         # Default behavior should catch the ValueError when trying to instantiate
         # the MVN and return NaN instead

test/optim/test_optimize.py

@@ -520,7 +520,7 @@ def test_optimize_acqf_nonlinear_constraints(self):
                     sequential=True,
                     raw_samples=16,
                 )
-            self.assertTrue(torch.allclose(candidates, 4 * torch.ones(3, **tkwargs)))
+            self.assertTrue(torch.allclose(candidates, 4 * torch.ones(1, 3, **tkwargs)))
 
             # Constrain the sum to be <= 4 in which case the solution is a
             # permutation of [4, 0, 0]
@@ -565,7 +565,9 @@ def nlc1(x):
                     batch_initial_conditions=batch_initial_conditions,
                     num_restarts=1,
                 )
-                self.assertTrue(torch.allclose(candidates, batch_initial_conditions))
+                self.assertTrue(
+                    torch.allclose(candidates, batch_initial_conditions[0, ...])
+                )
 
             # Constrain all variables to be >= 1. The global optimum is 2.45 and
             # is attained by some permutation of [1, 1, 2]
@@ -1248,7 +1250,7 @@ def test_optimize_acqf_discrete(self):
             )
             best_idcs = torch.topk(exp_acq_vals, q).indices
             expected_candidates = choices[best_idcs]
-            expected_acq_value = exp_acq_vals[best_idcs]
+            expected_acq_value = exp_acq_vals[best_idcs].reshape_as(acq_value)
             self.assertTrue(torch.allclose(acq_value, expected_acq_value))
             self.assertTrue(torch.allclose(candidates, expected_candidates))
 
@@ -1258,7 +1260,7 @@ def test_optimize_acqf_discrete(self):
             )
             best_idx = torch.argmax(exp_acq_vals)
             expected_candidates = choices[best_idx].repeat(q, 1)
-            expected_acq_value = exp_acq_vals[best_idx].repeat(q)
+            expected_acq_value = exp_acq_vals[best_idx].repeat(q).reshape_as(acq_value)
             self.assertTrue(torch.allclose(acq_value, expected_acq_value))
             self.assertTrue(torch.allclose(candidates, expected_candidates))
 
@@ -1268,7 +1270,7 @@ def test_optimize_acqf_discrete(self):
             )
             best_idcs = torch.topk(exp_acq_vals, q).indices
             expected_candidates = choices[best_idcs]
-            expected_acq_value = exp_acq_vals[best_idcs]
+            expected_acq_value = exp_acq_vals[best_idcs].reshape_as(acq_value)
             self.assertTrue(torch.allclose(acq_value, expected_acq_value))
             self.assertTrue(torch.allclose(candidates, expected_candidates))
 
@@ -1282,7 +1284,7 @@ def test_optimize_acqf_discrete(self):
             )
             best_idx = torch.argmax(exp_acq_vals)
             expected_candidates = choices[best_idx].repeat(q, 1)
-            expected_acq_value = exp_acq_vals[best_idx].repeat(q)
+            expected_acq_value = exp_acq_vals[best_idx].repeat(q).reshape_as(acq_value)
             self.assertTrue(torch.allclose(acq_value, expected_acq_value))
             self.assertTrue(torch.allclose(candidates, expected_candidates))
 

test/utils/multi_objective/box_decompositions/test_box_decomposition.py

@@ -318,7 +318,7 @@ def helper_hypervolume(self, Box_Decomp_cls: type) -> None:
             box_decomp = Box_Decomp_cls(ref_point=ref_point, Y=Y)
             hv = box_decomp.compute_hypervolume()
             self.assertEqual(hv.shape, expected_shape)
-            self.assertTrue(torch.allclose(hv, torch.tensor(0.0)))
+            self.assertTrue(torch.allclose(hv, torch.zeros(expected_shape)))
 
     def test_hypervolume(self) -> None:
         for cl in [