RewritePatternSet* patterns);

// Populates TF to MHLO legalization for some of the quantization ops.
//
// TODO(hinsu): Remove this once we combine quantized and non quantized op
// legalization in the ODML conversion pipeline.
void PopulateLegalizeTfQuantizationPatterns(MLIRContext* context,
                                            RewritePatternSet* patterns);

  AddShapeLegalizationPasses(pm);
  pm.addNestedPass<func::FuncOp>(
      CreateConvertCustomAggregationOpToQuantStatsPass());
  pm.addPass(createQuantizeCompositeFunctionsPass(options));
  // Add an inliner pass to inline quantized StableHLO functions.
  pm.addPass(createInlinerPass());
  if (pipeline_config.unpack_quantized_types()) {
    AddStablehloQuantToIntPasses(pm);
  }
}

void AddWeightOnlyQuantizationPasses(

limitations under the License.
==============================================================================*/
// This header file defines common utils used when transforming TF ops to
// Uniform Quantized ops.

#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_UTILS_TF_TO_UNIFORM_ATTRIBUTE_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_TENSORFLOW_UTILS_TF_TO_UNIFORM_ATTRIBUTE_UTILS_H_

absl::StatusOr<mlir::quant::QuantizedType> GetQuantizedType(
    const tflite::TensorT& tensor, mlir::Builder builder,
    bool is_constant = false, mlir::Type storage_type = {});

// Imports float tensor with calibration value into calibrated quantized type.
absl::StatusOr<mlir::quant::QuantizedType> GetCalibratedQuantizedType(
    const tflite::TensorT& tensor, mlir::Builder builder);

absl::StatusOr<mlir::TensorType> GetTensorType(const tflite::TensorT& tensor,

class UsedBy<string op> : Constraint<
  CPred<"llvm::isa<mlir::TFL::" # op # "Op>(*$0.getUsers().begin())">>;

// When the op is passing-through, the output types of the quantized ops need
// to be updated as well. Since the quantize op manages its own type by the
// "qtype" attribute, we should update the type shape in this attribute.
def ReorderTransposeDequantQuant :
      Pat<(TF_TransposeOp:$old_value

    "quantfork::QuantizationForkDialect"
  ];
  let options = [
      ListOption<"quantize_allowlist_", "quantize-allowlist", "std::string",
                 "comma separated list of allowlisted functions to be quantized. Only used in tests">,
      Option<"quantize_signed_", "quantize-signed", "bool", "false",
             "signed inference type. Only used in tests">,

  quantized_fraction = static_cast<int32_t>(std::round(fraction * (1L << 15)));
  // Clip extreme values.  These are more than enough to overflow int8, the
  // storage type for quantized values, and the final values will be clamped
  // no matter what.
  if (quantized_fraction == (1L << 15)) {
    quantized_fraction /= 2;
    ++shift;
  }
  if (shift < -15) {
    shift = 0;

namespace mlir::quant::stablehlo {

// A class that manages information about `QuantizableUnit`s post-quantization,
// internally in the form of `QuantizationUnits`. It is used to collect
// quantization summary from a quantized `ModuleOp` and emit it in a human- and
// machine-readable format.
class QuantizationReport {
 public:
  QuantizationReport() = default;

  // Initializes `QuantizationReport` by collecting `QuantizationResults` from

      "metadata": {
        "id": "cqeGmbO6PPNd"
      },
      "source": [
        "This example shows a JAX Keras reference model converted into a StableHLO module and via `jax2tf`, then quantized in the ODML Converter via the StableHLO Quantizer.\n",
        "\n",

// RUN: tfr-opt %s -tfr-decompose -tfr-rewrite-quantized-io -verify-diagnostics | FileCheck %s

// CHECK-LABEL: @tf__my_requantize
tfr.func @tf__my_requantize(%input: !tfr.tensor) -> !tfr.tensor {
  %raw_data = tfr.quant_raw_data(%input) : (!tfr.tensor) -> !tfr.tensor
  %scale, %zp = tfr.quant_qparam(%input) : (!tfr.tensor) -> (!tfr.tensor, !tfr.tensor)
  %result = tfr.call @tf__requantize(%raw_data, %scale, %zp) : (!tfr.tensor, !tfr.tensor, !tfr.tensor) -> !tfr.tensor