Avoid division by zero in normalize & unnormalize when lower & upper bounds are equal (#2363)

Summary:
Pull Request resolved: #2363

When lower and upper bounds are equal for a dimension, set `upper = lower + 1` to avoid division by zero in `normalize`. The behavior is mirrored in `unnormalize` to preserve `X = unnormalize(normalize(X, bounds), bounds)`.

This fixes some downstream issues observed in `get_chebyshev_objective`, and also lets us simplify `bounds` computation there.

Reviewed By: sdaulton

Differential Revision: D58168827

fbshipit-source-id: 37f29876fdab653c426a04b7c4a4122c8d41e497

Author: saitcakmak

botorch/utils/multi_objective/scalarization.py

@@ -95,14 +95,8 @@ def obj(Y: Tensor, X: Optional[Tensor] = None) -> Tensor:
             return -chebyshev_obj(Y=-Y)
 
         return obj
-    if Y.shape[-2] == 1:
-        # If there is only one observation, set the bounds to be
-        # [min(Y_m), min(Y_m) + 1] for each objective m. This ensures we do not
-        # divide by zero
-        Y_bounds = torch.cat([Y, Y + 1], dim=0)
-    else:
-        # Set the bounds to be [min(Y_m), max(Y_m)], for each objective m
-        Y_bounds = torch.stack([Y.min(dim=-2).values, Y.max(dim=-2).values])
+    # Set the bounds to be [min(Y_m), max(Y_m)], for each objective m.
+    Y_bounds = torch.stack([Y.min(dim=-2).values, Y.max(dim=-2).values])
 
     def obj(Y: Tensor, X: Optional[Tensor] = None) -> Tensor:
         # scale to [0,1]

botorch/utils/transforms.py

@@ -46,9 +46,32 @@ def standardize(Y: Tensor) -> Tensor:
     return (Y - Y.mean(dim=stddim, keepdim=True)) / Y_std
 
 
+def _update_constant_bounds(bounds: Tensor) -> Tensor:
+    r"""If the lower and upper bounds are identical for a dimension, set
+    the upper bound to lower bound + 1.
+
+    If any modification is needed, this will return a clone of the original
+    tensor to avoid in-place modification.
+
+    Args:
+        bounds: A `2 x d`-dim tensor of lower and upper bounds.
+
+    Returns:
+        A `2 x d`-dim tensor of updated lower and upper bounds.
+    """
+    if (constant_dims := (bounds[1] == bounds[0])).any():
+        bounds = bounds.clone()
+        bounds[1, constant_dims] = bounds[0, constant_dims] + 1
+    return bounds
+
+
 def normalize(X: Tensor, bounds: Tensor) -> Tensor:
     r"""Min-max normalize X w.r.t. the provided bounds.
 
+    NOTE: If the upper and lower bounds are identical for a dimension, that dimension
+    will not be scaled. Such dimensions will only be shifted as
+    `new_X[..., i] = X[..., i] - bounds[0, i]`. This avoids division by zero issues.
+
     Args:
         X: `... x d` tensor of data
         bounds: `2 x d` tensor of lower and upper bounds for each of the X's d
@@ -65,12 +88,17 @@ def normalize(X: Tensor, bounds: Tensor) -> Tensor:
         >>> bounds = torch.stack([torch.zeros(3), 0.5 * torch.ones(3)])
         >>> X_normalized = normalize(X, bounds)
     """
+    bounds = _update_constant_bounds(bounds=bounds)
     return (X - bounds[0]) / (bounds[1] - bounds[0])
 
 
 def unnormalize(X: Tensor, bounds: Tensor) -> Tensor:
     r"""Un-normalizes X w.r.t. the provided bounds.
 
+    NOTE: If the upper and lower bounds are identical for a dimension, that dimension
+    will not be scaled. Such dimensions will only be shifted as
+    `new_X[..., i] = X[..., i] + bounds[0, i]`, matching the behavior of `normalize`.
+
     Args:
         X: `... x d` tensor of data
         bounds: `2 x d` tensor of lower and upper bounds for each of the X's d
@@ -87,6 +115,7 @@ def unnormalize(X: Tensor, bounds: Tensor) -> Tensor:
         >>> bounds = torch.stack([torch.zeros(3), 0.5 * torch.ones(3)])
         >>> X = unnormalize(X_normalized, bounds)
     """
+    bounds = _update_constant_bounds(bounds=bounds)
     return X * (bounds[1] - bounds[0]) + bounds[0]
 
 

test/utils/multi_objective/test_scalarization.py

@@ -120,3 +120,13 @@ def test_get_chebyshev_scalarization(self):
                     + 0.05 * (weights * normalized_neg_Y_test).sum(dim=-1)
                 )
                 self.assertAllClose(Y_transformed, expected_Y_transformed)
+
+                # Test that it works when Y is constant in each dimension.
+                objective_transform = get_chebyshev_scalarization(
+                    weights=weights, Y=torch.zeros(2, 2, **tkwargs), alpha=0.0
+                )
+                Y_transformed = objective_transform(Y_test)
+                self.assertFalse(Y_transformed.isnan().any())
+                self.assertFalse(Y_transformed.isinf().any())
+                expected_Y = -(-weights * Y_test).max(dim=-1).values
+                self.assertAllClose(Y_transformed, expected_Y)

test/utils/test_transforms.py

@@ -59,7 +59,7 @@ def test_standardize(self):
 
 
 class TestNormalizeAndUnnormalize(BotorchTestCase):
-    def test_normalize_unnormalize(self):
+    def test_normalize_unnormalize(self) -> None:
         for dtype in (torch.float, torch.double):
             X = torch.tensor([0.0, 0.25, 0.5], device=self.device, dtype=dtype).view(
                 -1, 1
@@ -86,6 +86,17 @@ def test_normalize_unnormalize(self):
             self.assertTrue(torch.equal(X2_normalized, expected_X2_normalized))
             self.assertTrue(torch.equal(X2, unnormalize(X2_normalized, bounds=bounds2)))
 
+    def test_with_constant_bounds(self) -> None:
+        X = torch.rand(10, 2, dtype=torch.double)
+        # First dimension is constant, second has a range of 1.
+        # The transform should just add 1 to each dimension.
+        bounds = -torch.ones(2, 2, dtype=torch.double)
+        bounds[1, 1] = 0.0
+        X_normalized = normalize(X, bounds=bounds)
+        self.assertAllClose(X_normalized, X + 1)
+        X_unnormalized = unnormalize(X_normalized, bounds=bounds)
+        self.assertAllClose(X_unnormalized, X)
+
 
 class BMIMTestClass(BotorchTestCase):
     @t_batch_mode_transform(assert_output_shape=False)