size_of_splits = dim / splits;
      output_shape.push_back(size_of_splits);
    } else {
      output_shape.push_back(dim);
    }
  }

  SmallVector<mlir::Type, 4> output_types;
  for (int i = 0; i < splits; ++i) {
    output_types.push_back(
        mlir::RankedTensorType::get(output_shape, input_type.getElementType()));
  }
  auto size_of_splits_op = Create1DConstantOp(

      shape[split_dimension] = shape[split_dimension] / num_split;
      output_type =
          mlir::RankedTensorType::get(shape, input_type.getElementType());
    }
  } else {
    output_type = input_type;
  }

  // Creates a split op that splits |src_input| along |split_dimension|.
  llvm::SmallVector<mlir::Type, 4> output_types(num_split, output_type);
  *split_op = builder->create<mlir::TF::SplitOp>(

  %iter_handle1 = "tf.IteratorV2"() {container = "c", shared_name = "v0", output_shapes = [#tf_type.shape<>], output_types = [!tf_res]} : () -> !tf_res
  // expected-remark@below {{Result #0, ID 1 : 0, 1}}
  %iter_handle2 = "tf.IteratorV2"() {container = "c", shared_name = "v1", output_shapes = [#tf_type.shape<>], output_types = [!tf_res]} : () -> !tf_res
  func.return
}

// -----

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 {

      }

      // Collect all the quantized outputs and replace them by the results of
      // the new quantized op.
      llvm::SmallDenseMap<Value, int> outputs_replaced;
      SmallVector<Type, 4> output_types;
      output_types.reserve(quantizing_op->getNumResults());
      for (const auto& enumerated_result :
           llvm::enumerate(quantizing_op->getResults())) {
        Value result = enumerated_result.value();

    arg_types.push_back(arg.getType());
    arg_locs.push_back(arg.getLoc());
  }

  SmallVector<Type> output_types;
  for (const Value output : xla_call_module_op.getOutput()) {
    output_types.push_back(output.getType());
  }

  entry_func_op.setFunctionType(
      rewriter.getFunctionType(arg_types, output_types));

  // Replace argument types and locs.
  Block& entry = entry_func_op->getRegion(0).front();

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

                       ";\n");
  }
  if (op_info.output_types.empty()) {
    // Allow casting this class to Operation.
    strings::StrAppend(&class_decl,
                       "  operator ::tensorflow::Operation() const { "
                       "return operation; }\n");
  } else if (op_info.output_types.size() == 1) {
    if (op_info.is_list_output[0]) {

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 {