// `op_type` while the caller is the quantization unit.
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);

namespace {

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

class LiftQuantizableSpotsAsFunctionsPass

    OP_REQUIRES(ctx, tensor_data_writer_ != nullptr,
                absl::AbortedError("Could not create record writer"));

    // Fetch func_name and node_name from attributes and save as proto.
    quantization::UnitWiseQuantizationSpec::QuantizationUnit quant_unit_proto;
    quant_unit_proto.set_func_name(func_name);
    quant_unit_proto.set_node_name(node_name);

    string quant_unit_path = io::JoinPath(log_dir_path, "quant_unit.pb");

  TensorType tensor_type = 2;
}

// Unit-wise (op type and node name at this moment) quantization spec. It
// 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

#include "tensorflow/compiler/mlir/quantization/tensorflow/quantization_options.pb.h"

namespace mlir {
namespace quant {
namespace {

using QuantizationUnit =
    tensorflow::quantization::UnitWiseQuantizationSpec::QuantizationUnit;

// Adds QuantizationUnitLoc to quantizable layers.
class AddQuantizationUnitLocPass
    : public PassWrapper<AddQuantizationUnitLocPass,
                         OperationPass<func::FuncOp>> {

    'quantization.experimental.QuantizationComponentSpec'
)(quant_opts_pb2.QuantizationComponentSpec)

_UnitWiseQuantizationSpec = tf_export.tf_export(
    'quantization.experimental.UnitWiseQuantizationSpec'
)(quant_opts_pb2.UnitWiseQuantizationSpec)

_PresetMethod = _QuantizationMethod.PresetMethod
_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)

            quant_opts_pb2.UnitWiseQuantizationSpec(
                unit=[],
                quantization_method=quant_opts_pb2.QuantizationMethod(
                    preset_method=quant_opts_pb2.QuantizationMethod.METHOD_NO_QUANTIZE
                ),
            )
        ],
    )

    with self.assertRaisesRegex(
        ValueError, 'UnitWiseQuantizationSpec must contain at least one unit.'
    ):