// CHECK: return %[[RESULT]]
  // CHECK-NOT: tf.RecvTPUEmbeddingActivations
  // CHECK-NOT: tf.SendTPUEmbeddingGradients

  %0 = "tf.RecvTPUEmbeddingActivations"() {config = "\0A%\0A\0Dwatches_table\10\F5\03\18\80\02 \01*\0C\1A\00j\05\0D\00\00\80?\88\01\01\10\02\18\80\04 \01(\02"} : () -> tensor<512x256xf32>

  %2:2 = "tf.RecvTPUEmbeddingActivations"() {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D"} : () -> (tensor<2x2xf32>, tensor<4x4xf32>)
  func.return
}

// -----

// CHECK-LABEL: func @check_enqueue_ops_update_for_training

// RUN: tf-mlir-translate -mlir-to-graphdef %s -o - | FileCheck %s

func.func @main() {
  tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.RecvTPUEmbeddingActivations"() {config = "test_config_recv_embedding"} : () -> tensor<512x256xf32> loc("RecvTPUEmbedding")

    // CHECK-NEXT:   tf_device.return

    // CHECK:      "tf_device.cluster"
    // CHECK-NEXT:   "tf.RecvTPUEmbeddingActivations"
    // CHECK-NEXT:   "tf.SendTPUEmbeddingGradients"
    // CHECK-NEXT:   tf_device.return
    "tf_device.cluster"() ({
      %0 = "tf.RecvTPUEmbeddingActivations"() {config = "test_config_recv_embedding"} : () -> tensor<512x256xf32>

  auto output_ty =
      RankedTensorType::get({}, VariantType::get(region->getContext()));
  auto dedup_op = builder.create<XlaRecvTPUEmbeddingDeduplicationDataOp>(
      loc, output_ty, config);

  // Rewrite RecvTPUEmbeddingActivations op to the corresponding internal op.
  if (recv_op)
    AddOperandAndRewriteAs<XlaRecvTPUEmbeddingActivationsOp>(
        recv_op, dedup_op, recv_op->getAttrs(), &builder);

// CHECK-LABEL: func @ignore_embedding_ops
func.func @ignore_embedding_ops() -> () {
  "tf_device.cluster"() ({
    // CHECK: "tf.RecvTPUEmbeddingActivations"
    // CHECK-NOT: _xla_outside_compilation
    // CHECK: "tf.SendTPUEmbeddingGradients"
    // CHECK-NOT: _xla_outside_compilation
    %2:2 = "tf.RecvTPUEmbeddingActivations"() {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D"} : () -> (tensor<2x2xf32>, tensor<4x4xf32>)

results of these aggregations are visible to the Tensorflow Graph as the
outputs of a RecvTPUEmbeddingActivations op. This op returns a list containing
one Tensor of activations per table specified in the model. There can be at
most one RecvTPUEmbeddingActivations op in the TPU graph.
  }];

  let arguments = (ins
    StrAttr:$config
  );

  let results = (outs