auto output_fbm = tflite::FlatBufferModel::BuildFromBuffer(
      reinterpret_cast<const char*>(output_builder.GetCurrentBufferPointer()),
      output_builder.GetSize());
  tflite::ModelT output_model;
  output_fbm->GetModel()->UnPackTo(&output_model);

  // Extract output metadata
  std::map<std::string, std::string> output_metadata;
  for (const auto& metadata : output_model.metadata) {

Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))
    return block_arg.getOwner()->getParentOp();

  return value.getDefiningOp();
}

// Check if the TPUCopyWithDynamicShapeOp is valid.
// 1. The op should be wrapped inside a launch op.
// 2. The wrapped launch op should be placed on CPU.
LogicalResult CheckOpIsValid(Operation* op) {

  for (int non_bias : non_biases) {
    Operation *non_bias_define = op->getOperand(non_bias).getDefiningOp();
    if (auto dequant = llvm::dyn_cast<TFL::DequantizeOp>(non_bias_define)) {
      auto non_bias_type = mlir::cast<TensorType>(dequant.getInput().getType());
      auto non_bias_ele_type =
          mlir::cast<quant::QuantizedType>(non_bias_type.getElementType());
      non_bias_types.push_back(non_bias_ele_type);

                                PatternRewriter& rewriter) const override {
    if (call_op.getNumResults() != 1) return failure();
    auto users = call_op->getUsers();
    for (auto user : users) {
      if (!llvm::isa<mlir::stablehlo::UniformDequantizeOp>(user)) {
        return failure();
      }
    }
    auto func_name = call_op.getCallee();
    if (!func_name.starts_with("quantized_")) return failure();

      shape[0] = ShapedType::kDynamic;
      auto extensions =
          mhlo::TypeExtensionsAttr::get(func->getContext(), bounds);
      auto resultType =
          RankedTensorType::get(shape, inputType.getElementType(), extensions);
      arg.setType(resultType);
    }
    llvm::SmallVector<Type, 8> arg_types;
    for (auto arg : func.getArguments()) arg_types.push_back(arg.getType());
    func.setType(

                                                           int64_t* count) {
    auto weight = op->getOperand(1);
    auto weight_type =
        mlir::dyn_cast_or_null<mlir::RankedTensorType>(weight.getType());
    if (weight_type == nullptr || !weight_type.hasStaticShape()) return false;

    auto output = op->getResult(0);
    auto output_type =
        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());

  SmallVector<ElementsAttr> inputs;
  for (auto operand : op->getOperands()) {
    auto preceding_const_op = operand.getDefiningOp<TF::ConstOp>();
    if (preceding_const_op) {
      inputs.push_back(preceding_const_op.getValue());
      continue;
    }

    Operation* preceding_op = operand.getDefiningOp();
    int preceding_result_id = -1;
    for (auto preceding_result : preceding_op->getResults()) {

  auto function_uses =
      SymbolTable::getSymbolUses(target_func, &module_op.getBodyRegion());
  for (auto& function_use : function_uses.value()) {
    auto call_op = function_use.getUser();
    auto caller_func = call_op->getParentOfType<func::FuncOp>();
    if (!caller_func) return failure();

    builder.setInsertionPoint(call_op);
    for (auto [lifted_op, arg_idx] : lifted_op_and_arg_idx) {

  llvm::SmallVector<int32_t> result_shape_i32;
  auto result_shape = result_ty.getShape();
  for (auto element : result_shape) {
    result_shape_i32.push_back(static_cast<int32_t>(element));
  }
  auto result_shape_attr = builder.getI32TensorAttr(result_shape_i32);
  Value result_shape_tensor =
      builder.create<stablehlo::ConstantOp>(result_shape_attr);
  auto custom_code =

  // shaped tensors in the function.
  for (auto arg : func_op.getArguments()) {
    if (!mlir::isa<mlir::TensorType>(arg.getType())) continue;
    auto type = mlir::cast<mlir::TensorType>(arg.getType());
    if (type.hasRank()) {
      max_arg_size_ = std::max(max_arg_size_, GetRankedTensorSize(type));
    }
  }
}

void CostAnalysis::AnalyzeBlock(mlir::Block* block) {
  for (auto& op : *block) {
    EvaluateCost(&op);
  }