ArrayRef<int64_t> input_shape) {
  auto input_op = input.getDefiningOp();
  OpBuilder builder(input_op);
  builder.setInsertionPointAfter(input_op);
  SmallVector<int64_t, 4> transform_shape = {
      input_shape[0], input_shape[1] / block_size, input_shape[2] / block_size,
      input_shape[3] * block_size * block_size};
  auto transform_result_type =

    mlir::OpBuilder* builder) {
  // Inject concat ops after parallel_execute to merge outputs from
  // concurrently executed computations.
  builder->setInsertionPointAfter(new_parallel_execute);

  // Reorders outputs from TPUExecute op as defined by the output sharding
  // configuration.
  llvm::SmallVector<mlir::Value, 4> outputs_to_merge;

        }
        auto new_output_type = new_qtype.castFromExpressedType(
            mlir::quant::UniformQuantizedType::castToExpressedType(
                value.getType()));
        builder.setInsertionPointAfter(cst.getOperation());
        auto new_op = builder.create<tfl::QConstOp>(
            cst.getLoc(), new_output_type, mlir::TypeAttr::get(new_output_type),
            cst.getValueAttr());

    // Additionally, create the SymbolTable before further modifying the module.
    auto original_point = rewriter.saveInsertionPoint();
    rewriter.setInsertionPointAfter(op->getParentOfType<func::FuncOp>());
    SymbolTable manager(op->getParentOfType<ModuleOp>());

    // Constructs `then_branch`, which is executed when `if_cond` evaluates to
    // true.

      enqueue_callers = true;
    } else if (NeedsCastBack(use, tf_dialect_)) {
      if (!cast_op) {
        Operation* op = result.getDefiningOp();
        OpBuilder b(op);
        b.setInsertionPointAfter(op);
        cast_op = InsertCast(b, op->getLoc(), result.getType(), result);
        if (!cast_op) return false;
      }
      use.set(Value(cast_op->getResult(0)));
    }
  }

                                 control_nodes_at[owner]);
    }
    control_token.dropAllUses();

    // Replace the ControlNodeOp with the wrapped operation.
    rewriter.setInsertionPointAfter(outer_op);
    auto* cloned_inner = rewriter.clone(*inner_op);
    for (auto it :
         llvm::zip(control_node_op.getOutputs(), cloned_inner->getResults())) {