// checkpoint saving and restoring. This function returns a `SaverDef` instance
// with four fields populated: `version`, `filename_tensor_name`,
// `restore_op_name` and `save_tensor_name`. For valid quantized `graph_def` and
// `control_ret_node_names`, it should be able to retrieve the last three fields
// if there is at lest one variable in the graph.
//

    func.return %dot_out : tensor<*x!tf_type.qint32>
  }

  // Quantize initial input at the start of the graph. Output is qint8.
  func.func @quantize_i8(%input : tensor<*xf32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
    %quantize = "tf.UniformQuantize"(%input, %input_scale, %input_zp) {
      Tin = "tfdtype$DT_FLOAT",
      Tout = "tfdtype$DT_QINT8",

  auto func_op = module_op->lookupSymbol<func::FuncOp>("quantize");
  ASSERT_THAT(func_op, NotNull());

  auto uniform_quantize_op_itr =
      func_op.getBody().op_begin<mlir::stablehlo::UniformQuantizeOp>();
  ASSERT_THAT(
      uniform_quantize_op_itr,
      Ne(func_op.getBody().op_end<mlir::stablehlo::UniformQuantizeOp>()));

  // `uniform_quantize` is considered partially quantized because its output is

#include "mlir/IR/Builders.h"  // from @llvm-project

namespace mlir::quant {

// Caclulate padding values for XLA ops.
// Padding values for Uniform Quantized ops can be generated with this method as
// well as it shares the same definition for padding attribute with the XLA ops.
Value CalculatePaddingAndPadIfNeeded(OpBuilder &builder, Location loc,

// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions='quantization-method=weight_only target-opset=XLA' -quant-quantize-composite-functions='quantization-method=weight_only target-opset=XLA' -symbol-dce | FileCheck --check-prefix=PerTensor %s

          clEnumValN(OpSet::XLA, "XLA", "Uses TF XLA ops"),
          clEnumValN(OpSet::UNIFORM_QUANTIZED, "UNIFORM_QUANTIZED",
                     "Uses TF Uniform Quantized ops"))};

  Option<int64_t> min_num_elements_for_weights_{
      *this, "min-num-elements-for-weights", llvm::cl::init(0),
      llvm::cl::desc("The minimum required number of elements in a weight "

    auto complex_element_type = complex_type.getElementType();
    if (complex_element_type.isF32() || complex_element_type.isF64())
      return true;
  }
  // Check quantized types.
  if (auto quant_type = element_type.dyn_cast<mlir::quant::QuantizedType>()) {
    // TFLite supports QI16, QI32, QI8, and QUI8
    if ((quant_type.getStorageTypeIntegralWidth() == 16 &&
         quant_type.isSigned()) ||

      quantize_op.erase();
      arg.dropAllUses();
      bb.eraseArgument(0);
    };

    // This is looking for a pattern: arg -> tfl.quantize
    if (arg.hasOneUse() && llvm::isa<QuantizeOp>(*arg.user_begin())) {
      auto quantize_op = llvm::cast<QuantizeOp>(*arg.user_begin());
      remove_quantize_op(quantize_op);
      continue;
    }

Dialects can define entirely custom types, which is how MLIR can model things
like the LLVM IR type system (which has first class aggregates), domain
abstractions important for ML-optimized accelerators like quantized types, and
even the Swift or Clang type systems (which are built around Swift/Clang
declaration nodes) in the future.

If you want to connect a new low-level compiler, you would create a new dialect

      Constraint<CPred<"IsEinsumSupportedByXlaDotV2($0)">>;

// This attribute can be used in the `AttributeList` for missing attributes. It
// is necessary to keep other attributes in the same index as the quantized
// composite function.