void VisitOpOperands(GraphOp graph, Operation* op,
                     llvm::SmallPtrSetImpl<Operation*>* reachable_ops,
                     llvm::SmallVectorImpl<Operation*>* ops_to_visit) {
  for (Value operand : op->getOperands())
    VisitOpOperand(graph, operand, reachable_ops, ops_to_visit);
}

// Visits an op and it's associated operands. IslandOps are handled differently

  if (!return_op) return failure();
  if (return_op.getNumOperands() != 1) return failure();

  ReductionOp reduce_op = dyn_cast_or_null<ReductionOp>(
      return_op.getOperands().front().getDefiningOp());
  if (!reduce_op) return failure();
  if (reduce_op.getLhs() != body.getArgument(0) ||
      reduce_op.getRhs() != body.getArgument(1))
    return failure();

  return success();
}

      DCHECK_EQ(def->getParentOp(), graph);
      try_update_current_candidate(def);
    }

    // Check island data operands.
    island.walk([&](Operation* op) {
      for (Value input : op->getOperands()) {
        Operation* def = input.getDefiningOp();
        if (!def || def->getParentOp() != graph) continue;

        try_update_current_candidate(def);
      }
    });
  }

bool IsOpReplicateInvariant(Region* replicate_region, Operation* op) {
  auto ancestor_of_replicate = [&](Region* region) {
    return region && region->isProperAncestor(replicate_region);
  };

  for (Value operand : op->getOperands())
    if (!ancestor_of_replicate(operand.getParentRegion())) return false;

  // _TPUDeviceOrdinalPlaceholder implicitly depends on the replica.
  if (llvm::isa<TF::_TPUDeviceOrdinalPlaceholderOp>(op)) return false;

  llvm::SmallDenseMap<Value, Value> rewritten_operands;
  auto new_results =
      GetValueWithToken(builder, llvm::to_vector(terminator->getOperands()),
                        token, rewritten_operands, flatten_tuple);
  terminator->setOperands(new_results);
}

// Rewrites a `mhlo.if` op to receive and forward a `mhlo.token`. As If op does

  // Replace replicate terminator with YieldOp.
  builder.setInsertionPoint(&terminator);
  builder.create<tf_executor::YieldOp>(terminator.getLoc(),
                                       terminator.getOperands());
  terminator.erase();

  builder.setInsertionPoint(island_op);
  IRMapping mapping;
  for (int i : llvm::seq<int>(0, num_replicas)) {
    // Create new island for replica.

    yield.setOperand(0, identity.getOutput());
  } else {
    auto identity_n = builder.create<TF::IdentityNOp>(
        island.getLoc(), yield.getOperandTypes(), yield.getOperands());
    for (auto it : llvm::enumerate(identity_n.getResults()))
      yield.setOperand(it.index(), it.value());
  }
}

// Helper that creates an island. If `sub_op` is not nullptr, it will be moved

    cost_map_[op] = kDefaultCheapCost;
    return;
  }

  // For other ops, use the sum of input sizes as its cost.
  int64_t cost = kDefaultCheapCost;
  for (auto operand : op->getOperands()) {
    auto type = mlir::cast<mlir::TensorType>(operand.getType());
    if (type.hasRank()) {
      cost += GetRankedTensorSize(type);
    } else {

                                    OpBuilder* builder) {
  builder->create<func::ReturnOp>(cluster_return_op.getLoc(),
                                  cluster_return_op.getOperands());
  cluster_return_op.erase();
}

// Builds a function that outlines region attached to cluster_op or launch_op,
// and inserts built function into given module.
template <typename ClusterOrLaunchOp>

    yield.setOperand(0, identity.getOutput());
  } else {
    auto identity_n = builder.create<TF::IdentityNOp>(
        island.getLoc(), yield.getOperandTypes(), yield.getOperands());
    for (const auto& it : llvm::enumerate(identity_n.getResults()))
      yield.setOperand(it.index(), it.value());
  }
}

// Helper that creates an island that `sub_op` will be moved to.