Prevent unused NaNs from marking constraints as infeasible in pareto_…

…frontier_evaluator

Summary: `get_outcome_constraint_transforms` evaluates the constraints by taking the product of tensor `A` with `Y` and comparing the outcome to `rhs` (uses einsum for this). The product of `0` and `nan` evaluates to `nan`, leading to the constraint being marked infeasible, even if that `nan` is from some unused metric. By setting unused elements to `0`, this diff prevents such issues.

Differential Revision: D56953680

ax/models/tests/test_botorch_moo_defaults.py

@@ -183,6 +183,51 @@ def test_pareto_frontier_evaluator_with_outcome_constraints(self) -> None:
         )
         self.assertTrue(torch.equal(torch.tensor([2], dtype=torch.long), indx))
 
+    def test_pareto_frontier_evaluator_with_nan(self) -> None:
+        Y = torch.cat([self.Y, torch.zeros(5, 1)], dim=-1)
+        Yvar = torch.zeros(5, 4, 4)
+        weights = torch.tensor([1.0, 1.0, 0.0, 0.0])
+        outcome_constraints = (
+            torch.tensor([[0.0, 0.0, 1.0, 0.0]]),
+            torch.tensor([[3.5]]),
+        )
+        # Evaluate without NaNs as a baseline.
+        _, _, idx = pareto_frontier_evaluator(
+            model=None, objective_weights=weights, Y=Y, Yvar=Yvar
+        )
+        self.assertEqual(idx.tolist(), [2, 3, 4])
+        # Set an element of idx 2 to NaN. Should be removed.
+        Y[2, 1] = float("nan")
+        _, _, idx = pareto_frontier_evaluator(
+            model=None, objective_weights=weights, Y=Y, Yvar=Yvar
+        )
+        self.assertEqual(idx.tolist(), [3, 4])
+        # Set the unused constraint element of idx 3 to NaN. No effect.
+        Y[3, 2] = float("nan")
+        _, _, idx = pareto_frontier_evaluator(
+            model=None, objective_weights=weights, Y=Y, Yvar=Yvar
+        )
+        self.assertEqual(idx.tolist(), [3, 4])
+        # Add constraint, 3 should be removed.
+        _, _, idx = pareto_frontier_evaluator(
+            model=None,
+            objective_weights=weights,
+            Y=Y,
+            Yvar=Yvar,
+            outcome_constraints=outcome_constraints,
+        )
+        self.assertEqual(idx.tolist(), [4])
+        # Set unused index of 4 to NaN. No effect.
+        Y[4, 3] = float("nan")
+        _, _, idx = pareto_frontier_evaluator(
+            model=None,
+            objective_weights=weights,
+            Y=Y,
+            Yvar=Yvar,
+            outcome_constraints=outcome_constraints,
+        )
+        self.assertEqual(idx.tolist(), [4])
+
 
 class BotorchMOODefaultsTest(TestCase):
     def test_get_qLogEHVI_input_validation_errors(self) -> None:

ax/models/torch/botorch_moo_defaults.py

@@ -628,8 +628,13 @@ def pareto_frontier_evaluator(
     # Get feasible points that do not violate outcome_constraints
     if outcome_constraints is not None:
         cons_tfs = get_outcome_constraint_transforms(outcome_constraints)
+        # Handle NaNs in Y, if those elements are not part of the constraints.
+        # By setting the unused elements to 0, we prevent them from marking
+        # the whole constraint value as NaN and evaluating to infeasible.
+        Y_cons = Y.clone()
+        Y_cons[..., (outcome_constraints[0] == 0).all(dim=0)] = 0
         # pyre-ignore [16]
-        feas = torch.stack([c(Y) <= 0 for c in cons_tfs], dim=-1).all(dim=-1)
+        feas = torch.stack([c(Y_cons) <= 0 for c in cons_tfs], dim=-1).all(dim=-1)
         Y = Y[feas]
         Yvar = Yvar[feas]
         Y_obj = Y_obj[feas]