// inputs.
      SmallVector<Value, 4> inputs;
      inputs.reserve(op_with_region->getNumOperands());
      for (Value operand : op_with_region->getOperands()) {
        const Type operand_type = operand.getType();
        if (mlir::isa<NoneType>(operand_type)) {
          inputs.push_back(operand);
          continue;
        }

        const Type element_type =

    llvm::SetVector<Operation*>* unprocessed_subgraphs) {
  Subgraph& current_subgraph = subgraphs_.find(current_graph_call)->second;

  std::vector<Subgraph*> input_subgraphs;
  for (auto input : current_graph_call.getOperands()) {
    func::CallOp input_call = GetProducerCallOpOrNull(input);
    // If the input subgraph is not processed yet, we just go ahead and process
    // that one first.
    if (input_call == nullptr) continue;

OpFoldResult TFR::EqualOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  assert(operands.size() == 2 && "equal op has two operands");
  auto ctx = getContext();
  if (operands[0] == operands[1]) return BoolAttr::get(ctx, true);
  return BoolAttr::get(ctx, false);
}

OpFoldResult ConstOp::fold(FoldAdaptor adaptor) {
  auto operands = adaptor.getOperands();
  (void)operands;

    llvm::SmallVector<Value, 2> new_operands;
    FilterOutBlockArgControlDep(op->getOperands().drop_front(), new_operands);

    auto source_op = builder.create<tf_executor::NextIterationSourceOp>(
        op->getLoc(), op->getOperand(0).getType());
    builder.create<tf_executor::NextIterationSinkOp>(
        op->getLoc(), source_op.getToken(), /*input=*/op->getOperand(0),
        /*controlInputs=*/new_operands);
    op->replaceAllUsesWith(

bool HasDataDependencyWithUnscheduledOp(
    Operation& op, Block* block, SmallPtrSet<Operation*, 16>& unscheduled_ops) {
  WalkResult ready_to_schedule = op.walk([&](Operation* nested_op) {
    for (Value operand : nested_op->getOperands()) {
      Operation* defining_op = operand.getDefiningOp();
      if (!defining_op) continue;
      Operation* producer_in_block = block->findAncestorOpInBlock(*defining_op);

  }
  for (auto cluster_operand : cluster.getOperands()) {
    Operation* def = cluster_operand.getDefiningOp();
    // This pass assumes that a TPUPartitionedInputV2 is preceeded by
    // ReadVariable ops, and not vice versa. An earlier pass,
    // TPUResourceReadsWritesPartitioning, should have ensured this
    // precondition.
    if (!def) continue;
    for (auto operand : def->getOperands()) {

    Value operand = yield.getOperand(0);
    auto identity = builder.create<TF::IdentityOp>(island.getLoc(),
                                                   operand.getType(), operand);
    yield.setOperand(0, identity.getOutput());
  } else {
    auto identity_n = builder.create<TF::IdentityNOp>(
        island.getLoc(), yield.getOperandTypes(), yield.getOperands());

  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

    // `operand_idx`.
    return op.getNumOperands() + operand_idx;
  }
  return operand_idx;
}

// Returns the pointers to operands at `operand_indices` of `op`.
llvm::SmallVector<OpOperand*> GetOperands(Operation& op,
                                          llvm::ArrayRef<int> operand_indices) {
  llvm::SmallVector<OpOperand*> operands{};
  for (const int operand_idx : operand_indices) {

template <typename OpT>
void DeferRhsTransposeForBinaryOp(OpT op, PatternRewriter& rewriter) {
  auto transpose_op = cast<TransposeOp>(op.getOperand(0).getDefiningOp());
  Value lhs_pre_transpose = transpose_op.getOperand();

  // NCHW -> NHWC for the right-hand side, to match the operand's shape.
  Value rhs = op.getOperand(1);
  TransposeOp rhs_transpose_op = CreateTransposeOp(
      /*input=*/rhs, kNchwToNhwcPermutation, rewriter);