from tensorflow.python.platform import test

_CalibrationMethod = (
    stablehlo_quant_config_pb2.CalibrationOptions.CalibrationMethod
)


class CalibrationAlgorithmTest(test.TestCase, parameterized.TestCase):

  def test_min_max_max(self):
    calib_opts = stablehlo_quant_config_pb2.CalibrationOptions(
        calibration_method=_CalibrationMethod.CALIBRATION_METHOD_MIN_MAX
    )

        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",
        "//third_party/py/numpy",
    ],
)

pytype_strict_contrib_test(
    name = "calibration_algorithm_test",
    srcs = ["calibration_algorithm_test.py"],
    deps = [
        ":calibration_algorithm",
        ":calibration_statistics_proto_py",
        "//tensorflow/compiler/mlir/quantization/stablehlo:quantization_config_proto_py",