// 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;

// filled and node_name is always set to "_empty_node".
std::pair<std::string, std::string> GetFuncNameAndNodeName(
    TF::PartitionedCallOp call_op, const FlatSymbolRefAttr &f_attr) {
  std::optional<QuantizationUnitLoc::QuantizationUnit> quant_unit =
      FindQuantizationUnitFromLoc(call_op->getLoc());
  return std::make_pair(quant_unit->func_name(), quant_unit->node_name());
}

std::pair<std::string, std::string> GetFuncNameAndNodeName(

      return %0 : tensor<2x2xf32>
    }
    func.func private @composite_fn_1(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32> attributes {tf_quant.composite_function} {

  // Set the location of call op to QuantizationUnitLoc if found.
  Location call_op_loc = location;
  for (Operation* op : cloning_ops) {
    std::optional<QuantizationUnitLoc::QuantizationUnit> unit =
        FindQuantizationUnitFromLoc(op->getLoc());
    if (unit.has_value()) {
      call_op_loc = QuantizationUnitLoc(builder.getContext(), unit.value());
    }
  }

      # Verify if quant_unit.pb file was created correctly.
      quant_unit_file_path = os.path.join(log_dir_path, folder, 'quant_unit.pb')

      quant_unit = (
          quant_opts_pb2.UnitWiseQuantizationSpec.QuantizationUnit.FromString(
              open(quant_unit_file_path, 'rb').read()
          )
      )

      self.assertEqual(quant_unit.node_name, 'Conv2D')