Internal change.

PiperOrigin-RevId: 510569320
google · Feb 28, 2023 · 13fca22 · 13fca22
1 parent 2888c0f
commit 13fca22
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/vizier/_src/algorithms/designers/gp/acquisitions.py b/vizier/_src/algorithms/designers/gp/acquisitions.py
@@ -24,6 +24,7 @@
 from vizier.pyvizier import converters
 
 tfd = tfp.distributions
+tfp_bo = tfp.experimental.bayesopt
 
 
 class AcquisitionFunction(Protocol):
@@ -49,6 +50,18 @@ def __call__(
     return dist.mean() + 1.8 * dist.stddev()
 
 
+class EI(AcquisitionFunction):
+
+  def __call__(
+      self,
+      dist: tfd.Distribution,
+      features: Optional[chex.ArrayTree] = None,
+      labels: Optional[chex.Array] = None,
+  ) -> chex.Array:
+    del features
+    return tfp_bo.acquisition.GaussianProcessExpectedImprovement(dist, labels)()
+
+
 # TODO: Support discretes and categoricals.
 # TODO: Support custom distances.
 class TrustRegion: