Passes prior features to set-acquisition optimization

PiperOrigin-RevId: 634399230
google · May 16, 2024 · e1f2347 · e1f2347
1 parent 6da65e5
commit e1f2347
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/vizier/_src/algorithms/designers/gp_ucb_pe.py b/vizier/_src/algorithms/designers/gp_ucb_pe.py
@@ -826,15 +826,12 @@ def _suggest_one(
 
     if isinstance(acquisition_optimizer, vb.VectorizedOptimizer):
       acq_rng, self._rng = jax.random.split(self._rng)
-      prior_features = None
-      if self._all_completed_trials:
-        prior_features = vb.trials_to_sorted_array(
-            self._all_completed_trials, self._converter
-        )
       with profiler.timeit('acquisition_optimizer', also_log=True):
         best_candidates = eqx.filter_jit(acquisition_optimizer)(
             scoring_fn.score,
-            prior_features=prior_features,
+            prior_features=vb.trials_to_sorted_array(
+                self._all_completed_trials, self._converter
+            ),
             count=1,
             seed=acq_rng,
             score_with_aux_fn=scoring_fn.score_with_aux,
@@ -930,6 +927,9 @@ def _suggest_batch_with_exploration(
     with profiler.timeit('acquisition_optimizer', also_log=True):
       best_candidates = eqx.filter_jit(acquisition_optimizer)(
           scoring_fn.score,
+          prior_features=vb.trials_to_sorted_array(
+              self._all_completed_trials, self._converter
+          ),
           count=1,
           seed=acq_rng,
           score_with_aux_fn=scoring_fn.score_with_aux,