replicate_op->getLoc(), new_result_types, replicate_op->getOperands(),
      replicate_op->getAttrs());

  // Move region to the new op.
  new_replicate_op.getRegion().takeBody(replicate_op.getRegion());

  // Replace all old uses with new op results.
  int old_num_results = replicate_op->getNumResults();
  replicate_op->replaceAllUsesWith(

    int new_idx = child_idx >= cluster_idx ? child_idx + num_cores_per_replica
                                           : child_idx;
    new_parallel_execute->getRegions()[new_idx].takeBody(
        old_parallel_execute.getRegions()[old_idx]);
  }

  return cluster_idx;
}

mlir::tf_device::LaunchOp WrapOpInLaunch(mlir::OpBuilder* builder,
                                         mlir::Location loc,

      inner_op->setOperand(idx_args.index(), idx_args.value());
    }
    custom_op->setAttrs(inner_op->getAttrs());
    builder.create<YieldOp>(loc, inner_op->getResults());
    custom_op.getBody().takeBody(region);

    op->replaceAllUsesWith(custom_op);
    op->erase();
  }
}
}  // namespace

std::unique_ptr<OperationPass<func::FuncOp>> CreateRaiseCustomOpsPass() {

      parallel_execute.getLoc(), num_regions, new_result_types);

  for (auto region : llvm::zip(new_parallel_execute.getRegions(),
                               parallel_execute.getRegions()))
    std::get<0>(region).takeBody(std::get<1>(region));

  for (auto result :
       llvm::zip(results_to_remap, new_parallel_execute.getResults()))
    std::get<0>(result).replaceAllUsesWith(std::get<1>(result));

  parallel_execute.erase();

  tf_device::LaunchOp launch_op = builder->create<tf_device::LaunchOp>(
      builder->getUnknownLoc(), builder->getStringAttr(c.target),
      live_out_types);

  // Attach the region to launch_op.
  launch_op.getBody().takeBody(region);

  // Replace any external uses of live-out values with return values of launch
  // op. So live-out values no longer escape the region.
  ReplaceLiveOutExternalUses(live_outs, launch_op);

  auto new_launch_op = builder.create<tf_device::LaunchOp>(
      old_launch_op->getLoc(), old_launch_op->getDeviceAttr(),
      /*result_types=*/new_launch_op_results_types);

  new_launch_op.getBody().takeBody(old_launch_op->getBody());
  new_launch_op.GetBody().getTerminator()->setOperands(new_launch_op_results);

  return new_launch_op;
}

// Create the new device launch op which wraps the copy op.

    region.push_back(new_block);
    builder.setInsertionPointToEnd(&region.front());
    Operation *inner_op = builder.clone(*op);
    builder.create<YieldOp>(loc, inner_op->getResults());
    outer_op.getBody().takeBody(region);
    // Careful: We can't use outer_op.getResults(), because that also includes
    // the control token.
    op->replaceAllUsesWith(outer_op.getOutputs());
    op->erase();

    // We will "steal" the body of the island and replace it with a call to the
    // new function later.
    {
      YieldOp yield_op = island_op.GetYield();
      outlined_func.getBody().takeBody(island_op.getBody());

      // Replace the yield with a return
      OpBuilder replacer(yield_op);
      island_op.getBody().push_back(new Block);
      replacer.create<mlir::func::ReturnOp>(yield_op.getLoc(),