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")

//===----------------------------------------------------------------------===//
// The prepare-quantize-drq Pass.
//
namespace mlir {
namespace quant {

namespace {

using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;

// Applies prepare quantization on the model in TF dialect for dynamic range
// quantization case.
class PrepareQuantizeDRQPass

    %1 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsite("test@conv"("Conv2D_1") at "QuantizationUnit(\12\08Conv2D_1\1a\04conv)"))

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_dialect.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {
namespace {

using QuantizationUnit =
    ::tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;
using ::tensorflow::quantization::OpSet;
using ::tensorflow::quantization::QuantizationComponentSpec;
using ::tensorflow::quantization::QuantizationMethod;

//
namespace mlir {
namespace quant {

namespace {

using QuantMethod =
    ::tensorflow::quantization::QuantizationMethod::PresetMethod;
using QuantizationUnit = std::pair<Operation*, int>;
using QuantizationUnits = llvm::SetVector<QuantizationUnit>;
using ::tensorflow::quantization::OpSet;

// Preprocesses ops to allow multi-axis quantization, prior to quantization