const tfrt_stub::CostRecorder* cost_recorder = nullptr) {
    cost_recorder_ = cost_recorder;
    AnalyzeArguments(func_op);
    AnalyzeBlock(&func_op.front());
  }

  int64_t GetCost(mlir::Operation* op) const;

 private:
  void AnalyzeArguments(mlir::func::FuncOp func_op);
  void AnalyzeBlock(mlir::Block* block);
  void EvaluateCost(mlir::Operation* op);

  int64_t max_arg_size_ = 1;

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

void CostAnalysis::EvaluateCost(mlir::Operation* op) {
  if (auto cost_function =