// `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) {

      if (*this_partition == op_partition &&
          !unscheduled_op
               .walk([&](Operation *op) {
                 return llvm::all_of(op->getOperands(),
                                     [&](Value operand) {
                                       return is_ready(operand,
                                                       &unscheduled_op);

  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

  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();
}

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

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);

        func::ReturnOp return_op = *callee.getOps<func::ReturnOp>().begin();
        auto val = return_op.getOperand(i);
        auto block_arg_number = ComesFromBlockArgNumber(val);
        if (block_arg_number >= 0) {
          replaceAllUsesInRegionWith(call_op->getResult(i),
                                     call_op->getOperand(block_arg_number),
                                     func.getRegion());
        }
      }
    }

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;

  // 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.

      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);
      }
    });
  }