Fix hard-coded double precision in test_functions to default dtype (#2597)

Summary:
## Motivation

This PR replaces the hard-coded double precision that was used in the initialization of `test_functions/base.py` to use `torch.get_default_dtype()` instead.

See #2596 for more details.

### Have you read the [Contributing Guidelines on pull requests](https://github.com/pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Yes, I have read it.

Pull Request resolved: #2597

Test Plan:
I ran code formatting via `ufmt` and checked the code via `pytest -ra`. All tests related to the changes here were adjusted in the second commit of the branch.
Locally, two tests failed for me, but it seems to me that these are not related to the fix implemented here. If it turns out they are, I'd be more than happy to further adjust.

## Related PRs

None, but #2596 is related.

Reviewed By: saitcakmak

Differential Revision: D65066231

Pulled By: Balandat

fbshipit-source-id: 4beac1fc9a1e5094fd4806958ac2441a12506eb7

Author: AVHopp

botorch/test_functions/base.py

@@ -29,6 +29,7 @@ def __init__(
         self,
         noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""Base constructor for test functions.
 
@@ -37,6 +38,7 @@ def __init__(
                 provided, specifies separate noise standard deviations for each
                 objective in a multiobjective problem.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         super().__init__()
         self.noise_std = noise_std
@@ -47,7 +49,8 @@ def __init__(
                 f"Got {self.dim=} and {len(self._bounds)=}."
             )
         self.register_buffer(
-            "bounds", torch.tensor(self._bounds, dtype=torch.double).transpose(-1, -2)
+            "bounds",
+            torch.tensor(self._bounds, dtype=dtype).transpose(-1, -2),
         )
 
     def forward(self, X: Tensor, noise: bool = True) -> Tensor:
@@ -166,6 +169,7 @@ def __init__(
         self,
         noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""Base constructor for multi-objective test functions.
 
@@ -180,8 +184,8 @@ def __init__(
                 f"If specified as a list, length of noise_std ({len(noise_std)}) "
                 f"must match the number of objectives ({len(self._ref_point)})"
             )
-        super().__init__(noise_std=noise_std, negate=negate)
-        ref_point = torch.tensor(self._ref_point, dtype=torch.get_default_dtype())
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
+        ref_point = torch.tensor(self._ref_point, dtype=dtype)
         if negate:
             ref_point *= -1
         self.register_buffer("ref_point", ref_point)

botorch/test_functions/multi_fidelity.py

@@ -74,13 +74,19 @@ class AugmentedHartmann(SyntheticTestFunction):
     _optimizers = [(0.20169, 0.150011, 0.476874, 0.275332, 0.311652, 0.6573, 1.0)]
     _check_grad_at_opt = False
 
-    def __init__(self, noise_std: float | None = None, negate: bool = False) -> None:
+    def __init__(
+        self,
+        noise_std: float | None = None,
+        negate: bool = False,
+        dtype: torch.dtype = torch.double,
+    ) -> None:
         r"""
         Args:
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
         self.register_buffer("ALPHA", torch.tensor([1.0, 1.2, 3.0, 3.2]))
         A = [
             [10, 3, 17, 3.5, 1.7, 8],
@@ -126,13 +132,18 @@ class AugmentedRosenbrock(SyntheticTestFunction):
     _optimal_value = 0.0
 
     def __init__(
-        self, dim=3, noise_std: float | None = None, negate: bool = False
+        self,
+        dim=3,
+        noise_std: float | None = None,
+        negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             dim: The (input) dimension. Must be at least 3.
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         if dim < 3:
             raise ValueError(
@@ -141,7 +152,7 @@ def __init__(
         self.dim = dim
         self._bounds = [(-5.0, 10.0) for _ in range(self.dim)]
         self._optimizers = [tuple(1.0 for _ in range(self.dim))]
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
     def evaluate_true(self, X: Tensor) -> Tensor:
         X_curr = X[..., :-3]

botorch/test_functions/multi_objective.py

@@ -119,13 +119,15 @@ def __init__(
         self,
         noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the objectives.
+            dtype: The dtype that is used for the bounds of the function.
         """
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
         self._branin = Branin()
 
     def _rescaled_branin(self, X: Tensor) -> Tensor:
@@ -179,12 +181,14 @@ def __init__(
         dim: int,
         noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             dim: The (input) dimension.
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         if dim < self._min_dim:
             raise ValueError(f"dim must be >= {self._min_dim}, but got dim={dim}!")
@@ -194,7 +198,7 @@ def __init__(
         ]
         # max_hv is the area of the box minus the area of the curve formed by the PF.
         self._max_hv = self._ref_point[0] * self._ref_point[1] - self._area_under_curve
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
     @abstractmethod
     def _h(self, X: Tensor) -> Tensor:
@@ -339,13 +343,15 @@ def __init__(
         num_objectives: int = 2,
         noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             dim: The (input) dimension of the function.
             num_objectives: Must be less than dim.
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         if dim <= num_objectives:
             raise ValueError(
@@ -356,7 +362,7 @@ def __init__(
         self.k = self.dim - self.num_objectives + 1
         self._bounds = [(0.0, 1.0) for _ in range(self.dim)]
         self._ref_point = [self._ref_val for _ in range(num_objectives)]
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
 
 class DTLZ1(DTLZ):
@@ -608,12 +614,14 @@ def __init__(
         noise_std: None | float | list[float] = None,
         negate: bool = False,
         num_objectives: int = 2,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the objectives.
             num_objectives: The number of objectives.
+            dtype: The dtype that is used for the bounds of the function.
         """
         if num_objectives not in (2, 3, 4):
             raise UnsupportedError("GMM only currently supports 2 to 4 objectives.")
@@ -623,7 +631,7 @@ def __init__(
         if num_objectives > 3:
             self._ref_point.append(-0.1866)
         self.num_objectives = num_objectives
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
         gmm_pos = torch.tensor(
             [
                 [[0.2, 0.2], [0.8, 0.2], [0.5, 0.7]],
@@ -935,13 +943,15 @@ def __init__(
         num_objectives: int = 2,
         noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             dim: The (input) dimension of the function.
             num_objectives: Number of objectives. Must not be larger than dim.
             noise_std: Standard deviation of the observation noise.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         if num_objectives != 2:
             raise NotImplementedError(
@@ -954,7 +964,7 @@ def __init__(
         self.num_objectives = num_objectives
         self.dim = dim
         self._bounds = [(0.0, 1.0) for _ in range(self.dim)]
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
     @staticmethod
     def _g(X: Tensor) -> Tensor:
@@ -1246,15 +1256,17 @@ def __init__(
         noise_std: None | float | list[float] = None,
         constraint_noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             noise_std: Standard deviation of the observation noise of the objectives.
             constraint_noise_std: Standard deviation of the observation noise of the
                 constraint.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
         con_bounds = torch.tensor(self._con_bounds, dtype=self.bounds.dtype).transpose(
             -1, -2
         )
@@ -1357,6 +1369,7 @@ def __init__(
         noise_std: None | float | list[float] = None,
         constraint_noise_std: None | float | list[float] = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
@@ -1365,12 +1378,13 @@ def __init__(
             constraint_noise_std: Standard deviation of the observation noise of the
                 constraints.
             negate: If True, negate the function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         if dim < 2:
             raise ValueError("dim must be greater than or equal to 2.")
         self.dim = dim
         self._bounds = [(0.0, 1.0) for _ in range(self.dim)]
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
         self.constraint_noise_std = constraint_noise_std
 
     def LA2(self, A, B, C, D, theta) -> Tensor:

botorch/test_functions/sensitivity_analysis.py

@@ -24,13 +24,18 @@ class Ishigami(SyntheticTestFunction):
     """
 
     def __init__(
-        self, b: float = 0.1, noise_std: float | None = None, negate: bool = False
+        self,
+        b: float = 0.1,
+        noise_std: float | None = None,
+        negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             b: the b constant, should be 0.1 or 0.05.
             noise_std: Standard deviation of the observation noise.
             negative: If True, negative the objective.
+            dtype: The dtype that is used for the bounds of the function.
         """
         self._optimizers = None
         if b not in (0.1, 0.05):
@@ -52,7 +57,7 @@ def __init__(
             self.dgsm_gradient_square = [2.8, 24.5, 11]
         self._bounds = [(-math.pi, math.pi) for _ in range(self.dim)]
         self.b = b
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
     @property
     def _optimal_value(self) -> float:
@@ -127,13 +132,15 @@ def __init__(
         a: list = None,
         noise_std: float | None = None,
         negate: bool = False,
+        dtype: torch.dtype = torch.double,
     ) -> None:
         r"""
         Args:
             dim: Dimensionality of the problem. If 6, 8, or 15, will use standard a.
             a: a parameter, unless dim is 6, 8, or 15.
             noise_std: Standard deviation of observation noise.
-            negate: Return negatie of function.
+            negate: Return negative of function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         self._optimizers = None
         self.dim = dim
@@ -163,7 +170,7 @@ def __init__(
         else:
             self.a = a
         self.optimal_sobol_indicies()
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
     @property
     def _optimal_value(self) -> float:
@@ -207,11 +214,17 @@ class Morris(SyntheticTestFunction):
     Proposed to test sensitivity analysis methods
     """
 
-    def __init__(self, noise_std: float | None = None, negate: bool = False) -> None:
+    def __init__(
+        self,
+        noise_std: float | None = None,
+        negate: bool = False,
+        dtype: torch.dtype = torch.double,
+    ) -> None:
         r"""
         Args:
             noise_std: Standard deviation of observation noise.
             negate: Return negative of function.
+            dtype: The dtype that is used for the bounds of the function.
         """
         self._optimizers = None
         self.dim = 20
@@ -238,7 +251,7 @@ def __init__(self, noise_std: float | None = None, negate: bool = False) -> None
             0,
             0,
         ]
-        super().__init__(noise_std=noise_std, negate=negate)
+        super().__init__(noise_std=noise_std, negate=negate, dtype=dtype)
 
     @property
     def _optimal_value(self) -> float: