absl::flat_hash_map<int64_t, uint64_t> op_cost_map;
  op.walk([&](mlir::Operation* op) {
    const auto cost_attr = op->getAttrOfType<mlir::IntegerAttr>(kCostAttrName);
    if (!cost_attr) return;
    const auto op_key_attr =
        op->getAttrOfType<mlir::IntegerAttr>(kOpKeyAttrName);
    if (!op_key_attr) return;
    op_cost_map[op_key_attr.getInt()] = cost_attr.getInt();
  });
  return op_cost_map;
}

#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "tensorflow/core/platform/status.h"
#include "tensorflow/core/tfrt/fallback/cost_recorder.h"
#include "tensorflow/core/tfrt/fallback/op_cost_map.pb.h"

namespace tensorflow {
namespace tfrt_compiler {

// Analyze costs for tensorflow operations.
//
// The current heuristic used is quite simple, which is to calculate the total

                              op_with_dialect_count_[op_with_dialect_name]);
    }

    auto op_count_map = op_dtype_count_[op_with_dialect_name];
    if (!op_count_map.empty()) {
      std::string delim;
      *os_ << "  (";
      for (const auto &[dtype, cnt] : op_count_map) {
        *os_ << delim << absl::StrFormat("%s: %d", dtype, cnt);
        delim = ", ";
      }
      *os_ << ")";
    }