From 13fca22c7823287d314d9fe4b44cd5efd013e2c3 Mon Sep 17 00:00:00 2001
From: Emily Fertig <emilyaf@google.com>
Date: Fri, 17 Feb 2023 18:36:39 -0800
Subject: [PATCH] Internal change.

PiperOrigin-RevId: 510569320
---
 vizier/_src/algorithms/designers/gp/acquisitions.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/vizier/_src/algorithms/designers/gp/acquisitions.py b/vizier/_src/algorithms/designers/gp/acquisitions.py
index a1754e4cf..2ded9fd49 100644
--- 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: