Handle edge case when all labels are NaN.

PiperOrigin-RevId: 673447471

Author: xingyousong

vizier/_src/algorithms/designers/gp/output_warpers.py

@@ -455,11 +455,18 @@ class InfeasibleWarperComponent(OutputWarper):
   nearly equal to v_infeasible.
   """
 
-  _shift: float = attr.field(default=np.nan)
+  _shift: Optional[float] = attr.field(default=None)
 
   def warp(self, labels_arr: types.Array) -> types.Array:
     labels_arr = _validate_labels(labels_arr)
     labels_arr = labels_arr.flatten()
+
+    if np.isnan(labels_arr).all():
+      # Edge case when all values are NaN.
+      self._shift = np.nan
+      labels_arr[:] = 0
+      return labels_arr[:, np.newaxis]
+
     labels_range = np.nanmax(labels_arr) - np.nanmin(labels_arr)
     warped_bad_value = np.nanmin(labels_arr) - (0.5 * labels_range + 1)
     num_feasible = labels_arr.size - np.isnan(labels_arr).sum()
@@ -481,7 +488,7 @@ def warp(self, labels_arr: types.Array) -> types.Array:
     return labels_arr[:, np.newaxis]
 
   def unwarp(self, labels_arr: types.Array) -> types.Array:
-    if np.isnan(self._shift):
+    if self._shift is None:
       raise ValueError('warp() needs to be called before unwarp() is called.')
     return labels_arr - self._shift
 

vizier/_src/algorithms/designers/gp/output_warpers_test.py

@@ -490,6 +490,16 @@ def test_warper_removes_nans(self):
     labels_warped_infeasible = warper_infeasible.warp(labels)
     self.assertEqual(np.isnan(labels_warped_infeasible).sum(), 0)
 
+  def test_all_nans(self):
+    warper_infeasible = output_warpers.InfeasibleWarperComponent()
+    labels = np.array([[np.nan], [np.nan], [np.nan], [np.nan]])
+    labels_warped_infeasible = warper_infeasible.warp(labels)
+    expected = np.array([[0.0], [0.0], [0.0], [0.0]])
+    np.testing.assert_equal(labels_warped_infeasible, expected)
+
+    unwarped = warper_infeasible.unwarp(labels_warped_infeasible)
+    np.testing.assert_equal(unwarped, labels)
+
   def test_known_arrays(self):
     # TODO: Add a couple of parameterized test cases.
     self.skipTest('No test cases provided')