tf_device::ParallelExecuteOp BuildParallelExecuteOp(
    tf_device::ClusterFuncOp cluster_func, OpBuilder* builder) {
  const auto output_types = cluster_func.getResultTypes();
  builder->setInsertionPoint(cluster_func);
  auto parallel_execute = builder->create<tf_device::ParallelExecuteOp>(
      cluster_func.getLoc(), 1, output_types);
  cluster_func->remove();
  auto& block = parallel_execute.GetRegionBlockWithIndex(0);

  } : (tensor<!tf_type.variant>) -> tensor<!tf_type.variant>

  func.return %0 : tensor<!tf_type.variant>
// CHECK: "tf.MapDataset"(
// CHECK-SAME: <{f = @{{.*}}, metadata = "", output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string], preserve_cardinality = true, use_inter_op_parallelism = true}> {Targuments = []}
}

// CHECK-LABEL: TfTakeWhileDataset

func.func @main() {
  tf_executor.graph {
    %0 = tf_executor.island wraps "tf._TPUReplicate"() {computation = @foo, Tinputs = [], Tbroadcast_inputs = [], NumVariables = 0, Tguaranteed_constants = [], output_types = []} : () -> () loc("_TPUReplicate")
    tf_executor.fetch
  }
  func.return
}

func.func @foo() {
  tf_executor.graph {

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op
     (TF_MapDatasetOp $input_dataset, $other_arguments, $f, $output_types,
        $output_shapes, $use_inter_op_parallelism, $preserve_cardinality,
        $force_synchronous, $map_dataset_metadata),
     $batch_size, $drop_remainder, $parallel_copy, $batch_output_types,

    value {
      s: "name: \"batch_1\"\n[dist_belief.ImageInputDef.ext] {\n  num_rows: 128\n  num_cols: 128\n  mean_value: 128\n  std_value: 128\n  colorspace: RGB\n}\n"
    }
  }
  attr {
    key: "output_types"
    value {
      list {
        type: DT_FLOAT
      }
    }
  }
}
versions {
  producer: 27

func.func @main() {
  tf_executor.graph {
    %0 = tf_executor.island wraps "tf._TPUReplicate"() {computation = @foo, Tinputs = [], Tbroadcast_inputs = [], NumVariables = 0, Tguaranteed_constants = [], output_types = []} : () -> () loc("_TPUReplicate")
    tf_executor.fetch
  }
  func.return
}

func.func @foo() {
  tf_executor.graph {

      output_shapes = [#tf_type.shape<?>],
      output_types = [!tf_type.string]} : (tensor<*x!tf_type.variant>, tensor<i64>) -> tensor<!tf_type.variant>

    %1 = "tf.ReduceDataset"(%0, %cst_1, %arg0) {
      Targuments = [!tf_type.resource],
      Tstate = [i32], device = "",
      f = @__reduce_func, f._tf_data_function = true,
      output_shapes = [#tf_type.shape<>],

  %cst = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = ""} : () -> tensor<i32>
  %0 = "tf.MultiDeviceIteratorFromStringHandle"(%arg0) <{output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]}> {device = ""} : (tensor<!tf_type.string>) -> tensor<!tf_type.resource>

                            bool include_attr) const;

  string op_name;
  std::vector<string> arg_types;
  std::vector<string> arg_names;
  std::vector<string> output_types;
  std::vector<string> output_names;
  std::vector<bool> is_list_output;
  bool has_optional_attrs;
  string comment;

  const OpDef& graph_op_def;
  const ApiDef& api_def;

mlir::LogicalResult GetOutputTypesForLogicalDeviceComputation(
    int core_id, llvm::ArrayRef<xla::OpSharding> output_sharding_config,
    mlir::tf_device::ClusterFuncOp cluster_func,
    llvm::SmallVectorImpl<mlir::Type>* output_types,
    llvm::SmallVectorImpl<int>* cluster_to_core_index);

// Remaps outputs of `new_parallel_execute` op that represent concurrent
// execution of the `tf_device.cluster_func` at index `cluster_idx` of