tf_device.replicate([%arg0, %arg1] as %arg2: tensor<!tf_type.string>) {n = 2 : i32} {
        "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%arg2){table_ids = [1, 2]} : (tensor<!tf_type.string>) -> ()
        "tf.EnqueueTPUEmbeddingRaggedTensorBatch"(%arg2){table_ids = [1, 2]} : (tensor<!tf_type.string>) -> ()
        "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"(%arg2){table_ids = [1, 2]} : (tensor<!tf_type.string>) -> ()

      tf_device.replicate([%arg0, %arg1] as %arg2: tensor<!tf_type.string>) {n = 2 : i32} {
        "tf.EnqueueTPUEmbeddingSparseTensorBatch"(%arg2){table_ids = [1, 2]} : (tensor<!tf_type.string>) -> ()
        "tf.EnqueueTPUEmbeddingRaggedTensorBatch"(%arg2){table_ids = [1, 2]} : (tensor<!tf_type.string>) -> ()
        "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"(%arg2){table_ids = [1, 2]} : (tensor<!tf_type.string>) -> ()

        // expected-remark@above {{Successors: {4}}}
        "tf.EnqueueTPUEmbeddingRaggedTensorBatch"(%arg0) {table_ids = [1, 2], device_ordinal = 1} : (tensor<!tf_type.string>) -> ()
        // expected-remark@above {{ID: 0}}
        // expected-remark@above {{Successors: {1}}}
        "tf.EnqueueTPUEmbeddingRaggedTensorBatch"(%arg0) {table_ids = [1, 2], device_ordinal = 1} : (tensor<!tf_type.string>) -> ()

  func.func @embedding_head_extraction(%arg0: tensor<!tf_type.string>) {
    // CHECK:      "tf_device.launch"()
    // CHECK-SAME: device = "/job:worker/replica:0/task:0/device:CPU:0"
    // CHECK-NEXT:   "tf.EnqueueTPUEmbeddingRaggedTensorBatch"
    // CHECK-NEXT:   "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
    // CHECK-NEXT:   tf_device.return

    // CHECK:      "tf_device.cluster"
    // CHECK-NEXT:   "tf.UnknownOp"

  );

  let results = (outs);

  TF_DerivedOperandSizeAttr N = TF_DerivedOperandSizeAttr<0>;
}

def TF_EnqueueTPUEmbeddingRaggedTensorBatchOp : TF_Op<"EnqueueTPUEmbeddingRaggedTensorBatch", [DeclareOpInterfaceMethods<TF_GetResourceInstanceInterface>, SameVariadicOperandSize, TF_TPUEmbeddingWriteEffect]> {
  let summary = "Eases the porting of code that uses tf.nn.embedding_lookup().";