func.func @recv_tpu_embedding_activations() -> (tensor<512x256xf32>) {
  // CHECK: %[[DATA:.*]] = "tf.XlaRecvTPUEmbeddingDeduplicationData"() <{config = {{.*}}}> : () -> tensor<!tf_type.variant>
  // CHECK: %[[RESULT:.*]] = "tf.XlaRecvTPUEmbeddingActivations"(%[[DATA]]) <{config = {{.*}}}> : (tensor<!tf_type.variant>) -> tensor<512x256xf32>
  // CHECK: return %[[RESULT]]
  // CHECK-NOT: tf.RecvTPUEmbeddingActivations
  // CHECK-NOT: tf.SendTPUEmbeddingGradients

  let results = (outs
    TF_Tensor:$output
  );

  TF_DerivedResultTypeAttr Toutput = TF_DerivedResultTypeAttr<0>;
}

def TF_XlaRecvTPUEmbeddingActivationsOp : TF_Op<"XlaRecvTPUEmbeddingActivations", [TF_MustExecute, TF_TPUEmbeddingReadEffect]> {
  let summary = "An op that receives embedding activations on the TPU.";

  let description = [{