// of an effective cutoff value. In practice we could probably set it higher
// than what this function recommends.

var calibrate = flag.Bool("calibrate", false, "compute crossover for linear vs. binary search")

func TestCalibrate(t *testing.T) {
	if !*calibrate {
		return
	}

	if runtime.GOARCH == "amd64" {
		fmt.Printf("warning: running calibration on %s\n", runtime.GOARCH)
	}

// https://www.tensorflow.org/tutorials/load_data/tfrecord for details on the
// TFRecord format.
// Next ID: 2
message TfRecordFile {
  string path = 1;
}

// Configures a single representative dataset used to calibrate a single
// function.
// Next ID: 3
message RepresentativeDatasetConfig {
  oneof file {
    // Represents representative dataset saved as a .tfrecord file format.
    TfRecordFile tf_record = 1;
  }

// [Upper] [Lower] [Upper] [Lower].
const (
	UpperLower = MaxRune + 1 // (Cannot be a valid delta.)
)

// linearMax is the maximum size table for linear search for non-Latin1 rune.
// Derived by running 'go test -calibrate'.
const linearMax = 18

// is16 reports whether r is in the sorted slice of 16-bit ranges.
func is16(ranges []Range16, r uint16) bool {
	if len(ranges) <= linearMax || r <= MaxLatin1 {
		for i := range ranges {

from tensorflow.compiler.mlir.quantization.tensorflow import quantization_options_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_algorithm
from tensorflow.compiler.mlir.quantization.tensorflow.calibrator import calibration_statistics_pb2
from tensorflow.compiler.mlir.quantization.tensorflow.python import pywrap_function_lib

#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics.pb.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_average_min_max.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_histogram.h"
#include "tensorflow/compiler/mlir/quantization/tensorflow/calibrator/calibration_statistics_collector_min_max.h"

#include "tsl/platform/status.h"
#include "tsl/platform/status_matchers.h"

namespace tensorflow {
namespace {

using ::stablehlo::quantization::CalibrationOptions;
using ::tensorflow::calibrator::CalibrationStatistics;
using ::tensorflow::calibrator::CalibrationStatisticsMap;
using ::testing::Contains;
using ::testing::ElementsAre;
using ::testing::HasSubstr;
using ::testing::Key;
using ::testing::SizeIs;
using ::tsl::testing::StatusIs;

// CHECK: return %[[scale]], %[[zp]]
}

// CHECK-LABEL: quant_qparam_invalid
func.func @quant_qparam_invalid(%arg0: tensor<1x3x!quant.calibrated<f32<-1.0:1.0>>>) -> (!tfr.tensor, !tfr.tensor) {
  %0 = "tfr.cast"(%arg0) : (tensor<1x3x!quant.calibrated<f32<-1.0:1.0>>>) -> !tfr.tensor
  %scale, %zp = tfr.quant_qparam(%0) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
  func.return %scale, %zp: !tfr.tensor, !tfr.tensor

      // Quantize inputs of quantizable composite functions.
      for (OpOperand &input : op->getOpOperands()) {
        Type element_type = getElementTypeOrSelf(input.get().getType());
        // Non-float cases won't be calibrated.
        if (!element_type.isF32()) {
          continue;
        }

        // Skip when there is any already existing CustomAggregatorOp found.
        Operation *defining_op = input.get().getDefiningOp();

    %9, %9,
    %recurrent_stats, %cell_stats,
    %9, %9, %9, %9) {
      asymmetric_quantize_inputs = false,
      cell_clip = 1.000000e+01 : f32,
      effective_hidden_scale_intermediate = tensor<0x!quant.calibrated<f32<0.0:1.0>>>,
      fused_activation_function = "TANH",
      input_to_cell_intermediate = tensor<0xf32>,
      input_to_forget_intermediate = tensor<0xf32>,
      input_to_input_intermediate = tensor<0xf32>,

                                       elem_type.getIntOrFloatBitWidth());
  }

  // Intermediate tensors with calibration value (but not scale and zero points)
  // should return calibrated quantized type.
  if (is_intermediate && tensor.quantization != nullptr &&
      !IsQuantized(tensor)) {
    TF_ASSIGN_OR_RETURN(elem_type, GetCalibratedQuantizedType(tensor, builder));
  }