xla::XlaComputation* xla_computation, bool use_tuple_args,
    bool enable_op_fallback, bool return_tuple,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns =
        {},
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes = {},
    llvm::StringRef module_name = llvm::StringRef());

// Creates a MLIR pipeline that lowers MLIR module to MHLO dialect. The input

// true then the pipeline will include all the legalization passes.
void CreateConvertMlirToXlaHloPipeline(
    mlir::OpPassManager& pm, llvm::StringRef device_type,
    bool enable_op_fallback,
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes,
    bool lower_to_xla_hlo, bool allow_partial_conversion) {
  bool legalize_chlo = true;

  pm.addNestedPass<mlir::func::FuncOp>(

// Erase rewritten ClusterFuncOp(s). If TPUPartitionedInputV2Op /
// TPUPartitionedOutputV2Op are present, they must be removed along with the
// ClusterFuncOp(s).
mlir::LogicalResult EraseClusterFuncs(
    llvm::MutableArrayRef<mlir::tf_device::ClusterFuncOp> to_be_erased);

// Move child processes of the ParallelExecute that do not change. These are all
// children except for the child with the ClusterFunc.

    const TF::ResourceAliasAnalysis::Info& resource_alias_analysis) {
  // We need to know the devices to copy to.
  if (!replicate.getDevices()) return false;

  MutableArrayRef<OpOperand> inputs =
      replicate.GetOperandsForBlockArgument(replicate_arg);
  for (const auto& entry : llvm::enumerate(inputs)) {
    auto input_op = entry.value().get().getDefiningOp();

==============================================================================*/

#include "tensorflow/compiler/mlir/tensorflow/utils/xla_rewrite_util.h"

namespace tensorflow {
mlir::LogicalResult EraseClusterFuncs(
    llvm::MutableArrayRef<mlir::tf_device::ClusterFuncOp> to_be_erased) {
  for (auto cluster : to_be_erased) {
    auto old_parallel_execute =
        cluster->getParentOfType<mlir::tf_device::ParallelExecuteOp>();
    if (!old_parallel_execute) {

//
// Requires that block argument is of this replicate op.
MutableArrayRef<OpOperand> ReplicateOp::GetOperandsForBlockArgument(
    BlockArgument block_arg) {
  assert(block_arg.getOwner() == &GetBody());

  unsigned arg_number = block_arg.getArgNumber();
  unsigned num_replicated_args = GetNumReplicatedBlockArguments();
  int32_t num_replicas = getNAttr().getInt();
  MutableArrayRef<OpOperand> operands = getOperation()->getOpOperands();

Status CompileMlirToXlaHloViaBuilder(
    mlir::ModuleOp module_op, llvm::ArrayRef<TensorOrResourceShape> arg_shapes,
    llvm::StringRef device_type, XlaCompilationResult* compilation_result,
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes) {
  // This call to RefineShapes is redundant with the call in BuildHloFromTf.
  // It's here so xla::Parameters that are created form block.getArguments will

  // of indices from the original values to the target values.
  void TransposeRecursively(const ArrayRef<float> original_values,
                            const MutableArrayRef<float> target_values,
                            SmallVector<int64_t>& current_indices) const {
    // Map an element from `original_values` to `target_values` when a set of
    // indices is formed.

    unsigned GetReplicaOperandIndexForBlockArgument(BlockArgument block_arg, unsigned replica);
    Value GetReplicaOperandForBlockArgument(BlockArgument block_arg, unsigned replica);
    MutableArrayRef<OpOperand> GetOperandsForBlockArgument(BlockArgument block_arg);
    bool WrapsSingleOp();
  }];

  let builders = [
    OpBuilder<(ins "int":$n,

  return new_island;
}

// Creates respective YieldOp for the new merged island.
YieldOp CreateNewIslandYieldOp(IslandOp new_island,
                               llvm::MutableArrayRef<IslandResult> results) {
  llvm::SmallVector<Value, 8> yield_operands;
  yield_operands.reserve(results.size());

  for (auto [old_result, new_island] :
       llvm::zip(results, new_island.getOutputs())) {