// CHECK: "tf.Iterator"
  // CHECK-SAME: output_shapes{{.*}}200x28x28x1{{.*}}200x10
  // CHECK-SAME: output_types = [f32, f32]
  // CHECK-SAME: shared_name = "_iterator1"
  %0 = "tf.AnonymousIteratorV3"() {output_shapes = [
    #tf_type.shape<200x28x28x1>,
    #tf_type.shape<200x10>], output_types = [f32, f32]} : () -> tensor<!tf_type.resource>
  // CHECK: "tf.Iterator"
  // CHECK-SAME: shared_name = "_iterator2"

  auto indices_type = RankedTensorType::get(input_shape, builder->getI64Type());
  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
  SmallVector<mlir::Type, 2> input_types{input_type, indices_type};
  SmallVector<mlir::Type, 1> output_types{output_type};
  return createMaxUnpoolingFunc<2, 1>(builder, input_types, output_types);
}

template <int N>

    // CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "funcB"} : 1

    // CHECK: tfrt_fallback_async.executeop key({{[0-9]+}}) cost({{.*}}) device("/device:CPU:0") "tf.FlatMapDataset"({{.*}}) {Targuments = [], metadata = "", output_shapes = [#corert.shape<>], output_types = [i64]} {f = "__inference_Dataset_flat_map_lambda_19"} : 1

    %1 = "tf.ReduceDataset"(%arg0, %arg1) {
      Targuments = [],
      Tstate = [i64], device = "",
      f = @__reduce_func_0, f._tf_data_function = true,
      output_shapes = [#tf_type.shape<>],
      output_types = [i64], use_inter_op_parallelism = true } : (tensor<!tf_type.variant>, tensor<i64>) -> (tensor<i64>)
    func.return
}

           output_shapes = [#tf_type.shape<>], output_types = [i32],
           preserve_cardinality = false, sloppy = false,
           use_inter_op_parallelism = true,
           metadata = ""} : (tensor<*x!tf_type.variant>) -> tensor<!tf_type.variant>
  %5 = "tf.BatchDatasetV2"(%4, %0, %1) {device = "",
           output_shapes = [#tf_type.shape<>],
           output_types = [i32],
           parallel_copy = false,

  computeInputsOutputs(ops, &inputs, &outputs);

  std::vector<Type> input_types;
  std::vector<Type> output_types;

  input_types.reserve(inputs.size());
  for (Value v : inputs) {
    input_types.push_back(v.getType());
  }
  output_types.reserve(outputs.size());
  for (Value v : outputs) {
    output_types.push_back(v.getType());
  }

  // Create the function.
  MLIRContext* context = ops[0]->getContext();

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

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,