class QuantizationUnitLoc : public CallSiteLoc {
 public:
  using QuantizationUnit =
      tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(QuantizationUnitLoc)

  QuantizationUnitLoc(MLIRContext* context, const QuantizationUnit& unit);

  // Checks if the given location is QuantizationUnitLoc. Users could call

#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace quant {
namespace {

// Prefix and suffix to the QuantizationUnit string representation.
constexpr std::string_view kQuantizationUnitPrefix = "QuantizationUnit(";
constexpr std::string_view kQuantizationUnitSuffix = ")";

// Concatenates node name and func name with a "@" separator.

  } else {
    for (Location child_loc : locations) {
      FindQuantizationUnitsRecursively(child_loc, units);
    }
  }
}

// Finds the QuantizationUnit from location.
std::optional<QuantizationUnit> FindQuantizationUnit(Operation* op) {
  SmallVector<QuantizationUnit> quant_units;
  FindQuantizationUnitsRecursively(op->getLoc(), quant_units);

  if (quant_units.size() == 1) {
    return *quant_units.begin();
  }

  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
  return %3 : tensor<1x3x2x2xf32>
// CHECK: tf.Conv2D
// CHECK-SAME: loc(callsite("Model/conv2d@conv2d_with_valid_loc"("Conv2D") at "QuantizationUnit({{.*}})"))
}

func.func @conv2d_with_callsite_loc(%arg0: tensor<1x3x4x3xf32>) -> (tensor<1x3x2x2xf32>) {

// CHECK-SAME: _tfl_quant_trait = "fully_quantizable"
// CHECK-SAME: loc(callsite("Model/conv2d@conv2d_unmatching_unit"("Conv2D") at "QuantizationUnit({{.*}})"))

// -----

module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1269 : i32}} {

}

// `DumpTensor` op saves entire value of input to as a tensor proto into a
// specified directory and filename. When enabled is set to false, op is
// disabled and won't save any value. It also creates `QuantizationUnit` proto
// with `func_name` and `node_name` to identify the op.
REGISTER_OP("DumpTensor")
    .Input("tensor_data: T")
    .Attr("log_dir_path: string")
    .Attr("file_name: string")
    .Attr("T: type")

// contains the scope of the unit, and the quantization method for each unit.
// NEXT ID: 7
message UnitWiseQuantizationSpec {
  // Quantization unit granularity.
  // NEXT ID: 4
  message QuantizationUnit {
    // Type of the op, ex: Conv2D, MatMul, Einsum... The node_name field can
    // be omitted if it is intended to match all nodes with this type.
    string op_type = 1;