Review comments applied.

tensorflow/python/tpu/gpu_embedding.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Mid level API for TPU Embeddings without Embedding Accelerator."""
+"""GPU embeddings API."""
 
 from typing import Any, Dict, Iterable, Optional, Text, Union
 
@@ -37,14 +37,13 @@
 
 @tf_export("tpu.experimental.embedding.GPUEmbedding")
 class GPUEmbeddingV0(tpu_embedding_v1.TPUEmbeddingV0):
-  """The GPUEmbedding mid level API running on TPU without Embedding accelerator.
+  """GPU embeddings API.
 
-  NOTE: This mid level API is not intended for large embedding table lookup.
+  NOTE: This API is not intended for large embedding table lookup.
   Embedding tables will be replicated across devices rather than sharding
   across them. To do large embedding table lookup, please use the
   `tpu.experimental.embedding.GPUEmbedding` class. This class is an alternative
-  way to do embedding lookups when the TPU doesn't support any version of
-  embedding feature. See
+  way to do embedding lookups on GPUs. See
   `tpu.experimental.tpu_hardware_feature.embedding_feature` for a detailed
   explanation.
 

tensorflow/python/tpu/tpu_hardware_feature.py

@@ -42,6 +42,8 @@ class EmbeddingFeature(enum.Enum):
     V2: Embedding lookup accelerator V2. The embedding lookup operation can be
         placed anywhere of the computation. Multiple instances of embedding
         lookup layer is allowed.
+    GPU:Embedding lookup on GPU. No embedding lookup accelerator features are
+        used.
     """
     UNSUPPORTED = "UNSUPPORTED"
     V1 = "V1"