Clean up float32 defaults in Sobol utils & test functions (#2346)

Summary:
Pull Request resolved: #2346

We recommend using double/float64 in BoTorch but seem to default to float32 in some Sobol utilities and test function code. I searched through the code base and updated float32 defaults in a few ways:
- Either removed them and left them to use whatever default torch sees fit.
- Set them to `torch.get_default_dtype`, to explicitly match the dtype used by torch.
- Updated them to use the dtype of another tensor from the same class.

Some added `dtype=torch.get_default_dtype()` is unnecessarily verbose, as torch will default to it whenever the input is `List[Union[int, float]]` or `List[float]` -- but it will default to `torch.long` if the input is `List[int]`. I think being verbose here helps avoid confusion.

Found while investigating #2264.

Also made some minor updates in these files to make pyre happier (or less upset, whichever fits).

Reviewed By: Balandat

Differential Revision: D57473132

fbshipit-source-id: cf98b792f8f34abf20d8e4b7bbf714fdc73b764e

Author: saitcakmak

botorch/sampling/qmc.py

@@ -58,7 +58,10 @@ def __init__(
         self._sobol_engine = SobolEngine(dimension=sobol_dim, scramble=True, seed=seed)
 
     def draw(
-        self, n: int = 1, out: Optional[Tensor] = None, dtype: torch.dtype = torch.float
+        self,
+        n: int = 1,
+        out: Optional[Tensor] = None,
+        dtype: Optional[torch.dtype] = None,
     ) -> Optional[Tensor]:
         r"""Draw `n` qMC samples from the standard Normal.
 
@@ -67,10 +70,12 @@ def draw(
             out: An option output tensor. If provided, draws are put into this
                 tensor, and the function returns None.
             dtype: The desired torch data type (ignored if `out` is provided).
+                If None, uses `torch.get_default_dtype()`.
 
         Returns:
             A `n x d` tensor of samples if `out=None` and `None` otherwise.
         """
+        dtype = torch.get_default_dtype() if dtype is None else dtype
         # get base samples
         samples = self._sobol_engine.draw(n, dtype=dtype)
         if self._inv_transform:

botorch/test_functions/base.py

@@ -11,10 +11,9 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import List, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
-
 from botorch.exceptions.errors import InputDataError
 from torch import Tensor
 from torch.nn import Module
@@ -146,16 +145,16 @@ def evaluate_slack_true(self, X: Tensor) -> Tensor:
         pass  # pragma: no cover
 
 
-class MultiObjectiveTestProblem(BaseTestProblem):
-    r"""Base class for test multi-objective test functions.
+class MultiObjectiveTestProblem(BaseTestProblem, ABC):
+    r"""Base class for multi-objective test functions.
 
     TODO: add a pareto distance function that returns the distance
     between a provided point and the closest point on the true pareto front.
     """
 
     num_objectives: int
     _ref_point: List[float]
-    _max_hv: float
+    _max_hv: Optional[float] = None
 
     def __init__(
         self,
@@ -176,16 +175,16 @@ def __init__(
                 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.float)
+        ref_point = torch.tensor(self._ref_point, dtype=torch.get_default_dtype())
         if negate:
             ref_point *= -1
         self.register_buffer("ref_point", ref_point)
 
     @property
     def max_hv(self) -> float:
-        try:
+        if self._max_hv is not None:
             return self._max_hv
-        except AttributeError:
+        else:
             raise NotImplementedError(
                 f"Problem {self.__class__.__name__} does not specify maximal "
                 "hypervolume."

botorch/test_functions/multi_fidelity.py

@@ -46,12 +46,12 @@ class AugmentedBranin(SyntheticTestFunction):
     def evaluate_true(self, X: Tensor) -> Tensor:
         t1 = (
             X[..., 1]
-            - (5.1 / (4 * math.pi**2) - 0.1 * (1 - X[:, 2])) * X[:, 0] ** 2
+            - (5.1 / (4 * math.pi**2) - 0.1 * (1 - X[:, 2])) * X[:, 0].pow(2)
             + 5 / math.pi * X[..., 0]
             - 6
         )
         t2 = 10 * (1 - 1 / (8 * math.pi)) * torch.cos(X[..., 0])
-        return t1**2 + t2 + 10
+        return t1.pow(2) + t2 + 10
 
 
 class AugmentedHartmann(SyntheticTestFunction):
@@ -93,15 +93,15 @@ def __init__(self, noise_std: Optional[float] = None, negate: bool = False) -> N
             [1312, 1696, 5569, 124, 8283, 5886],
             [2329, 4135, 8307, 3736, 1004, 9991],
             [2348, 1451, 3522, 2883, 3047, 6650],
-            [4047, 8828, 8732, 5743, 1091, 381],
+            [4047, 8828, 8732, 5743, 1091, 381.0],
         ]
-        self.register_buffer("A", torch.tensor(A, dtype=torch.float))
-        self.register_buffer("P", torch.tensor(P, dtype=torch.float))
+        self.register_buffer("A", torch.tensor(A))
+        self.register_buffer("P", torch.tensor(P))
 
     def evaluate_true(self, X: Tensor) -> Tensor:
         self.to(device=X.device, dtype=X.dtype)
         inner_sum = torch.sum(
-            self.A * (X[..., :6].unsqueeze(-2) - 0.0001 * self.P) ** 2, dim=-1
+            self.A * (X[..., :6].unsqueeze(-2) - 0.0001 * self.P).pow(2), dim=-1
         )
         alpha1 = self.ALPHA[0] - 0.1 * (1 - X[..., 6])
         H = (
@@ -147,6 +147,6 @@ def __init__(
     def evaluate_true(self, X: Tensor) -> Tensor:
         X_curr = X[..., :-3]
         X_next = X[..., 1:-2]
-        t1 = 100 * (X_next - X_curr**2 + 0.1 * (1 - X[..., -2:-1])) ** 2
-        t2 = (X_curr - 1 + 0.1 * (1 - X[..., -1:]) ** 2) ** 2
+        t1 = 100 * (X_next - X_curr.pow(2) + 0.1 * (1 - X[..., -2:-1])).pow(2)
+        t2 = (X_curr - 1 + 0.1 * (1 - X[..., -1:]).pow(2)).pow(2)
         return -((t1 + t2).sum(dim=-1))

botorch/test_functions/multi_objective.py

@@ -76,7 +76,7 @@
 import math
 from abc import ABC, abstractmethod
 from math import pi
-from typing import List, Union
+from typing import List, Tuple, Union
 
 import torch
 from botorch.exceptions.errors import UnsupportedError
@@ -170,7 +170,7 @@ class DH(MultiObjectiveTestProblem, ABC):
     """
 
     num_objectives = 2
-    _ref_point: float = [1.1, 1.1]
+    _ref_point: List[float] = [1.1, 1.1]
     _x_1_lb: float
     _area_under_curve: float
     _min_dim: int
@@ -286,7 +286,7 @@ class DH3(DH):
     _min_dim = 3
 
     @staticmethod
-    def _exp_args(x: Tensor) -> Tensor:
+    def _exp_args(x: Tensor) -> Tuple[Tensor, Tensor]:
         exp_arg_1 = -((x - 0.35) / 0.25).pow(2)
         exp_arg_2 = -((x - 0.85) / 0.03).pow(2)
         return exp_arg_1, exp_arg_2
@@ -1078,7 +1078,7 @@ def gen_pareto_front(self, n: int) -> Tensor:
                 torch.linspace(
                     left + self._eps,
                     right - self._eps,
-                    n_per_part[i],
+                    int(n_per_part[i]),
                     dtype=self.bounds.dtype,
                     device=self.bounds.device,
                 )
@@ -1195,12 +1195,12 @@ class BNH(MultiObjectiveTestProblem, ConstrainedBaseTestProblem):
 
     def evaluate_true(self, X: Tensor) -> Tensor:
         return torch.stack(
-            [4.0 * (X**2).sum(dim=-1), ((X - 5.0) ** 2).sum(dim=-1)], dim=-1
+            [4.0 * X.pow(2).sum(dim=-1), (X - 5.0).pow(2).sum(dim=-1)], dim=-1
         )
 
     def evaluate_slack_true(self, X: Tensor) -> Tensor:
-        c1 = 25.0 - (X[..., 0] - 5.0) ** 2 - X[..., 1] ** 2
-        c2 = (X[..., 0] - 8.0) ** 2 + (X[..., 1] + 3.0) ** 2 - 7.7
+        c1 = 25.0 - (X[..., 0] - 5.0).pow(2) - X[..., 1].pow(2)
+        c2 = (X[..., 0] - 8.0).pow(2) + (X[..., 1] + 3.0).pow(2) - 7.7
         return torch.stack([c1, c2], dim=-1)
 
 
@@ -1256,7 +1256,9 @@ def __init__(
             negate: If True, negate the function.
         """
         super().__init__(noise_std=noise_std, negate=negate)
-        con_bounds = torch.tensor(self._con_bounds, dtype=torch.float).transpose(-1, -2)
+        con_bounds = torch.tensor(self._con_bounds, dtype=self.bounds.dtype).transpose(
+            -1, -2
+        )
         self.register_buffer("con_bounds", con_bounds)
         self.constraint_noise_std = constraint_noise_std
 
@@ -1373,7 +1375,7 @@ def __init__(
         super().__init__(noise_std=noise_std, negate=negate)
         self.constraint_noise_std = constraint_noise_std
 
-    def LA2(self, A, B, C, D, theta):
+    def LA2(self, A, B, C, D, theta) -> Tensor:
         return A * torch.sin(B * theta.pow(C)).pow(D)
 
     def evaluate_true(self, X: Tensor) -> Tensor:
@@ -1420,22 +1422,22 @@ class OSY(MultiObjectiveTestProblem, ConstrainedBaseTestProblem):
 
     def evaluate_true(self, X: Tensor) -> Tensor:
         f1 = -(
-            25 * (X[..., 0] - 2) ** 2
-            + (X[..., 1] - 2) ** 2
-            + (X[..., 2] - 1) ** 2
-            + (X[..., 3] - 4) ** 2
-            + (X[..., 4] - 1) ** 2
+            25 * (X[..., 0] - 2).pow(2)
+            + (X[..., 1] - 2).pow(2)
+            + (X[..., 2] - 1).pow(2)
+            + (X[..., 3] - 4).pow(2)
+            + (X[..., 4] - 1).pow(2)
         )
-        f2 = (X**2).sum(-1)
+        f2 = X.pow(2).sum(-1)
         return torch.stack([f1, f2], dim=-1)
 
     def evaluate_slack_true(self, X: Tensor) -> Tensor:
         g1 = X[..., 0] + X[..., 1] - 2.0
         g2 = 6.0 - X[..., 0] - X[..., 1]
         g3 = 2.0 - X[..., 1] + X[..., 0]
         g4 = 2.0 - X[..., 0] + 3.0 * X[..., 1]
-        g5 = 4.0 - (X[..., 2] - 3.0) ** 2 - X[..., 3]
-        g6 = (X[..., 4] - 3.0) ** 2 + X[..., 5] - 4.0
+        g5 = 4.0 - (X[..., 2] - 3.0).pow(2) - X[..., 3]
+        g6 = (X[..., 4] - 3.0).pow(2) + X[..., 5] - 4.0
         return torch.stack([g1, g2, g3, g4, g5, g6], dim=-1)
 
 
@@ -1453,12 +1455,12 @@ class SRN(MultiObjectiveTestProblem, ConstrainedBaseTestProblem):
     _ref_point = [0.0, 0.0]  # TODO: Determine proper reference point
 
     def evaluate_true(self, X: Tensor) -> Tensor:
-        obj1 = 2.0 + ((X - 2.0) ** 2).sum(dim=-1)
-        obj2 = 9.0 * X[..., 0] - (X[..., 1] - 1.0) ** 2
+        obj1 = 2.0 + (X - 2.0).pow(2).sum(dim=-1)
+        obj2 = 9.0 * X[..., 0] - (X[..., 1] - 1.0).pow(2)
         return torch.stack([obj1, obj2], dim=-1)
 
     def evaluate_slack_true(self, X: Tensor) -> Tensor:
-        c1 = 225.0 - ((X**2) ** 2).sum(dim=-1)
+        c1 = 225.0 - (X.pow(2)).pow(2).sum(dim=-1)
         c2 = -10.0 - X[..., 0] + 3 * X[..., 1]
         return torch.stack([c1, c2], dim=-1)
 
@@ -1490,8 +1492,8 @@ def evaluate_true(self, X: Tensor) -> Tensor:
         # different numbers (see below).
         # f1 = WeldedBeam.evaluate_true(self, X)
         x1, x2, x3, x4 = X.unbind(-1)
-        f1 = 1.10471 * (x1**2) * x2 + 0.04811 * x3 * x4 * (14.0 + x2)
-        f2 = 2.1952 / (x4 * x3**3)
+        f1 = 1.10471 * x1.pow(2) * x2 + 0.04811 * x3 * x4 * (14.0 + x2)
+        f2 = 2.1952 / (x4 * x3.pow(3))
         return torch.stack([f1, f2], dim=-1)
 
     def evaluate_slack_true(self, X: Tensor) -> Tensor:
@@ -1509,17 +1511,17 @@ def evaluate_slack_true(self, X: Tensor) -> Tensor:
         # g3 = 1 / (5 - 0.125) * g3_
         # g4 = 1 / P * g6_
 
-        R = torch.sqrt(0.25 * (x2**2 + (x1 + x3) ** 2))
+        R = torch.sqrt(0.25 * (x2.pow(2) + (x1 + x3).pow(2)))
         M = P * (L + x2 / 2)
         # This `J` is different than the one in [CoelloCoello2002constraint]_
         # by a factor of 2 (sqrt(2) instead of sqrt(0.5))
-        J = 2 * math.sqrt(0.5) * x1 * x2 * (x2**2 / 12 + 0.25 * (x1 + x3) ** 2)
+        J = 2 * math.sqrt(0.5) * x1 * x2 * (x2.pow(2) / 12 + 0.25 * (x1 + x3).pow(2))
         t1 = P / (math.sqrt(2) * x1 * x2)
         t2 = M * R / J
-        t = torch.sqrt(t1**2 + t1 * t2 * x2 / R + t2**2)
-        s = 6 * P * L / (x4 * x3**2)
+        t = torch.sqrt(t1.pow(2) + t1 * t2 * x2 / R + t2.pow(2))
+        s = 6 * P * L / (x4 * x3.pow(2))
         # These numbers are also different from [CoelloCoello2002constraint]_
-        P_c = 64746.022 * (1 - 0.0282346 * x3) * x3 * x4**3
+        P_c = 64746.022 * (1 - 0.0282346 * x3) * x3 * x4.pow(3)
 
         g1 = (t - t_max) / t_max
         g2 = (s - s_max) / s_max