template <typename OpT>
OpT AddOperandAndRewriteAs(Operation* op, Value operand, NamedAttrList attr,
                           OpBuilder* builder) {
  builder->setInsertionPoint(op);
  auto operands = llvm::to_vector<4>(op->getOperands());
  operands.push_back(operand);
  auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
                                     operands, attr.getAttrs());

                                   SmallVector<TransposeOp, 2>* transpose_ops) {
  for (auto it = transpose_ops->begin(); it != transpose_ops->end(); ++it) {
    auto tranpose_op = *it;
    for (auto tranpose_operand : tranpose_op.getOperands()) {
      auto ranked_tranpose_type =
          mlir::dyn_cast_or_null<RankedTensorType>(tranpose_operand.getType());
      if (!ranked_tranpose_type) continue;
      if (ranked_tranpose_type.getRank() == permutation.size() &&

        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();
    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
        return_op.getOperands()[0].getDefiningOp());
    if (!req_op) return failure();

    // Create a new func.call op with f32 output.
    auto new_call_op = call_op.clone();
    new_call_op->getResult(0).setType(

  branch_types.reserve(branches.size());

  if (llvm::any_of(op->getOperands(),
                   [](Value value) { return value == nullptr; }))
    return op->emitOpError("operation has null operand");

  // Functions have one less operand compared to op as first operand is elided
  // (`cond` of `tf.If` and `branch_index` of `tf.Case`).
  TypeRangeWithDesc input{op->getOperands().drop_front().getTypes(), "input"};

                                            tf_device::LaunchOp launch) {
  // Forward launch inner op results to launch op results.
  launch.replaceAllUsesWith(launch.GetBody().getTerminator()->getOperands());

  // For all inner ops, assign the launch device as a `device` attribute.
  if (failed(AssignDevicesInRegion(tf_dialect, launch, launch.getBody())))
    return failure();

    }

    if (read_operands.empty()) continue;

    // Update caller and function types with new read operands.
    auto operands = llvm::to_vector<4>(cluster_func.getOperands());
    operands.append(read_operands.begin(), read_operands.end());

    auto loc = cluster_func.getLoc();
    auto new_cluster_func = builder.create<tf_device::ClusterFuncOp>(

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()) {

      // original If op.
      if (op_to_inline.hasTrait<OpTrait::IsTerminator>()) {
        updated_results.reserve(op_to_inline.getNumOperands());
        for (Value operand : op_to_inline.getOperands())
          updated_results.push_back(mapper.lookup(operand));
        break;
      }

      // Otherwise, clone the op here.
      rewriter.clone(op_to_inline, mapper);
    }

  block.walk([&](Operation *op) {
    if (op->getName().getStringRef() != "tfg.NextIteration") return;
    mlir::OpBuilder builder(op);

    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>(