Generate test for NdarrayOptimizationClosure that surfaces failures that was otherwise flaky (#1505)

Summary:
Pull Request resolved: #1505

We have been seeing
```Caused by ValueError: failed to initialize intent(inout) array -- expected elsize=8 but got 4``` on occasion. Although this should be fixed by D41355824 (c7ed6ab), this was a frustratingly flaky error, happening ~1%-4% of the time in unit tests, so we wanted a test that errors consistently to guard against regressions. The error happens as a result of NaNs being cast to the wrong dtype upon optimization faiures in `NdarrayOptimizationClosure`. We didn't want to iterate on the flaky test, `test_fit.TestFitOther.test_fit_with_converter` since it will likely be delted soon, so I added a new test that creates this issue more directly.

Reviewed By: j-wilson

Differential Revision: D41386524

fbshipit-source-id: cc23cfd6857c924c240b6fcb33edd58827e384cd

Author: esantorella

test/optim/test_fit.py

@@ -142,6 +142,34 @@ def _test_fit_gpytorch_mll_scipy(self, mll):
                 )
             mock_closure.assert_called_once_with(ab="cd")
 
+    def test_fit_with_nans(self) -> None:
+        """Test the branch of NdarrayOptimizationClosure that handles errors."""
+
+        from botorch.optim.closures import NdarrayOptimizationClosure
+
+        def closure():
+            raise RuntimeError("singular")
+
+        for dtype in [torch.float32, torch.float64]:
+
+            parameters = {"x": torch.tensor([0.0], dtype=dtype)}
+
+            wrapper = NdarrayOptimizationClosure(closure=closure, parameters=parameters)
+
+            def _assert_np_array_is_float64_type(array) -> bool:
+                # e.g. "float32" in "torch.float32"
+                self.assertEqual(str(array.dtype), "float64")
+
+            _assert_np_array_is_float64_type(wrapper()[0])
+            _assert_np_array_is_float64_type(wrapper()[1])
+            _assert_np_array_is_float64_type(wrapper.state)
+            _assert_np_array_is_float64_type(wrapper._get_gradient_ndarray())
+
+            # Any mll will do
+            mll = next(iter(self.mlls.values()))
+            # will error if dtypes are wrong
+            fit.fit_gpytorch_mll_scipy(mll, closure=wrapper, parameters=parameters)
+
 
 class TestFitGPyTorchMLLTorch(BotorchTestCase):
     def setUp(self):

test/test_fit.py

@@ -642,26 +642,34 @@ def mock_fit_gpytorch_mll(*args, **kwargs):
 
 
 class TestFitOther(BotorchTestCase):
-    def test_fit_with_converter(self):
+    def helper_fit_with_converter(self, dtype) -> None:
         # Check that sequential optimization using converter does not
         # break input transforms.
-        for dtype in (torch.float, torch.double):
-            tkwargs = {"device": self.device, "dtype": dtype}
-            X = torch.rand(5, 2, **tkwargs) * 10
-            Y = X**2
-            intf = Normalize(2)
-            model = SingleTaskGP(X, Y, input_transform=intf)
-            mll = ExactMarginalLogLikelihood(model.likelihood, model)
-            with patch(
-                f"{fit_gpytorch_mll.__module__}.batched_to_model_list",
-                wraps=batched_to_model_list,
-            ) as wrapped_converter, warnings.catch_warnings(record=True) as ws:
-                warnings.simplefilter("always", BotorchWarning)
-                fit_gpytorch_mll(mll)
-            # Check that MLL repacking succeeded.
-            self.assertFalse(
-                any("Training loss of repacked model" in str(w.message) for w in ws)
-            )
-            wrapped_converter.assert_called_once()
-            self.assertFalse(torch.allclose(intf.mins, torch.zeros(1, 2, **tkwargs)))
-            self.assertFalse(torch.allclose(intf.ranges, torch.ones(1, 2, **tkwargs)))
+        tkwargs = {"device": self.device, "dtype": dtype}
+        # Set the seed to a number that doesn't generate numerical
+        # issues (no NaNs)
+        torch.manual_seed(0)
+        X = torch.rand(5, 2, **tkwargs) * 10
+        Y = X**2
+        intf = Normalize(2)
+        model = SingleTaskGP(X, Y, input_transform=intf)
+        mll = ExactMarginalLogLikelihood(model.likelihood, model)
+        with patch(
+            f"{fit_gpytorch_mll.__module__}.batched_to_model_list",
+            wraps=batched_to_model_list,
+        ) as wrapped_converter, warnings.catch_warnings(record=True) as ws:
+            warnings.simplefilter("always", BotorchWarning)
+            fit_gpytorch_mll(mll)
+        # Check that MLL repacking succeeded.
+        self.assertFalse(
+            any("Training loss of repacked model" in str(w.message) for w in ws)
+        )
+        wrapped_converter.assert_called_once()
+        self.assertFalse(torch.allclose(intf.mins, torch.zeros(1, 2, **tkwargs)))
+        self.assertFalse(torch.allclose(intf.ranges, torch.ones(1, 2, **tkwargs)))
+
+    def test_fit_with_converter_float32(self) -> None:
+        self.helper_fit_with_converter(torch.float)
+
+    def test_fit_with_converter_float64(self) -> None:
+        self.helper_fit_with_converter(torch.double)