Trial and metric padding is not yet supported for multimetric problem…

…s in GP-UCB-PE. PiperOrigin-RevId: 716793840
google · Jan 17, 2025 · af6f150 · af6f150
1 parent 5ab515f
commit af6f150
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 13 deletions.
diff --git a/vizier/_src/algorithms/designers/gp_ucb_pe.py b/vizier/_src/algorithms/designers/gp_ucb_pe.py
@@ -655,14 +655,24 @@ def __attrs_post_init__(self):
     # Extra validations
     if self._problem.search_space.is_conditional:
       raise ValueError(f'{type(self)} does not support conditional search.')
-    elif (
-        len(self._problem.metric_information) != 1
-        and self._config.optimize_set_acquisition_for_exploration
-    ):
-      raise ValueError(
-          f'{type(self)} works with exactly one metric when'
-          ' `optimize_set_acquisition_for_exploration` is enabled.'
-      )
+    elif len(self._problem.metric_information) != 1:
+      if self._config.optimize_set_acquisition_for_exploration:
+        raise ValueError(
+            f'{type(self)} works with exactly one metric when'
+            ' `optimize_set_acquisition_for_exploration` is enabled.'
+        )
+      empty_labels = jnp.array([[]])
+      padded_labels = self._padding_schedule.pad_labels(empty_labels)
+      if padded_labels.shape[0] != empty_labels.shape[0]:
+        raise ValueError(
+            f'{type(self)} does not support trial padding for multimetric'
+            ' problems.'
+        )
+      if padded_labels.shape[1] != empty_labels.shape[1]:
+        raise ValueError(
+            f'{type(self)} does not support metric padding for multimetric'
+            ' problems.'
+        )
 
     # Extra initializations.
     # Discrete parameters are continuified to account for their actual values.

diff --git a/vizier/_src/algorithms/designers/gp_ucb_pe_test.py b/vizier/_src/algorithms/designers/gp_ucb_pe_test.py
@@ -90,7 +90,6 @@ class GpUcbPeTest(parameterized.TestCase):
           iters=3,
           batch_size=3,
           num_metrics=2,
-          applies_padding=True,
           multimetric_promising_region_penalty_type=(
               gp_ucb_pe.MultimetricPromisingRegionPenaltyType.UNION
           ),
@@ -99,7 +98,6 @@ class GpUcbPeTest(parameterized.TestCase):
           iters=3,
           batch_size=3,
           num_metrics=2,
-          applies_padding=True,
           ensemble_size=4,
           multimetric_promising_region_penalty_type=(
               gp_ucb_pe.MultimetricPromisingRegionPenaltyType.INTERSECTION
@@ -111,23 +109,20 @@ class GpUcbPeTest(parameterized.TestCase):
           num_seed_trials=5,
           num_metrics=2,
           multitask_type=mt_type.SEPARABLE_NORMAL_TASK_KERNEL_PRIOR,
-          applies_padding=True,
       ),
       dict(
           iters=3,
           batch_size=5,
           num_seed_trials=5,
           num_metrics=2,
           multitask_type=mt_type.SEPARABLE_LKJ_TASK_KERNEL_PRIOR,
-          applies_padding=True,
       ),
       dict(
           iters=3,
           batch_size=5,
           num_seed_trials=5,
           num_metrics=2,
           multitask_type=mt_type.SEPARABLE_DIAG_TASK_KERNEL_PRIOR,
-          applies_padding=True,
       ),
   )
   def test_on_flat_space(