}

void collect_input_defining_islands(IslandOp op, std::vector<IslandOp>& ops) {
  Operation* wrapped_op = &op.GetBody().front();
  for (Value operand : wrapped_op->getOperands()) {
    IslandOp wrapper = dyn_cast_or_null<IslandOp>(operand.getDefiningOp());
    if (!wrapper || !wrapper.WrapsSingleOp()) continue;
    ops.push_back(wrapper);
  }
}

void collect_output_users_islands(IslandOp op, std::vector<IslandOp>& ops) {

  // Construct a new island from the given root.
  explicit MergedIsland(IslandOp root) : insert_point(root) {
    islands.push_back(root);
  }

  // The insertion point anchor of the merged island, or where the merged island
  // will be inserted when created.
  Operation* const insert_point;

  // The set of islands that are to be merged together.
  SmallVector<IslandOp> islands;
};

  // We only call this function for island or fetch ops. The second pair of
  // parentheses is needed for successful compilation.
  assert((isa<IslandOp, FetchOp>(op)));
  mlir::MutableOperandRange mutable_control_inputs =
      isa<IslandOp>(op) ? cast<IslandOp>(op).getControlInputsMutable()
                        : cast<FetchOp>(op).getFetchesMutable();
  // Add control inputs in program order of the defining ops.

// (from YieldOps first operand) is forwarded to the IslandOp single data result
// users.
struct DropEmptyIslandNoOperandOneDataResult
    : public OpRewritePattern<IslandOp> {
  using OpRewritePattern<IslandOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(IslandOp op,
                                PatternRewriter &rewriter) const override {

    new_op->setAttrs(item.getAttrDictionary());
    item.erase();
  }
  AddStatefulControlFlowDependencies(graph_op);
}

// Populates an empty IslandOp and with a NoOp or Identity/IdentityN depending
// on if there are any data results.
void PopulateEmptyIsland(tf_executor::IslandOp island) {
  OpBuilder builder(&island.GetBody(), island.GetBody().begin());
  tf_executor::YieldOp yield = island.GetYield();

                           ShapedType const_type) {
  IslandOp while_wrapper = while_op->getParentOfType<IslandOp>();
  assert(while_wrapper && "While op is expected to be wrapped in a IslandOp");

  // Get the dummy constant.
  OpBuilder builder(while_wrapper);
  auto loc = NameLoc::get(
      builder.getStringAttr("chain_control_outputs@" + while_op.getBody()));
  IslandOp const_wrapper = GetDummyConstant(builder, const_type, loc);

    cloned_fetch_op.erase();
  };

  return llvm::to_vector(cloned_fetch_op.getFetches());
}

// Creates a new `IslandOp` that wraps a `TF::NoOp`. The `IslandOp` has control
// dependencies to the values provided.
IslandOp CreateNoOpWithControlDependencies(
    const Location loc, GraphOp main_graph_op,
    const ArrayRef<Value> control_dependencies) {

// for the associated replica island.
LogicalResult ExpandReplicateIntoReplicas(
    const Dialect* tf_dialect, OpBuilder& builder,
    tf_executor::IslandOp island_op, tf_device::ReplicateOp replicate_op,
    int num_replicas, llvm::SmallVectorImpl<tf_executor::IslandOp>& replicas,
    bool legacy_graph_export, int replica_group_idx) {
  replicas.reserve(num_replicas);
  auto devices = replicate_op.getDevices();

  cloned_fetch_op.erase();

  return control_fetch;
}

// Creates a new `IdentityOp` wrapped by an `IslandOp`. The identity op returns
// the `main_file_prefix_arg` and has control dependencies to `control_inputs`.
IslandOp CreateFilePrefixIdentityOp(const BlockArgument main_file_prefix_arg,
                                    const ArrayRef<Value> control_inputs,

    return ops_have_same_device(sink.GetSource(), sink);
  }

  return llvm::isa<tf_executor::EnterOp, tf_executor::ExitOp,
                   tf_executor::IslandOp, tf_executor::MergeOp,
                   tf_executor::SwitchOp>(op);
}

// Assigns all data results to a specified device.
void PopulateDeviceForOpResults(
    Operation& op, llvm::StringRef device,