op_set=target_opset,
    )

    if target_opset != quant_opts_pb2.XLA:
      # Uniform quantized opset is not supported for weight-only
      with self.assertRaisesRegex(
          ValueError, 'TF/Uniform quantized opset does not support weight-only.'
      ):
        converted_model = quantize_model.quantize(
            input_saved_model_path,
            output_directory,
            quantization_options,

    Runs the post-calibration passes for post-training quantization.
  }];
  let options = [
    Option<"unpack_quantized_types_", "unpack-quantized-types", "bool",
      /*default=*/"true", "Unpacks ops with uniform quantized types into "
      "operations without uniform quantized types (mostly i8 or i32).">
  ];
  let dependentDialects = [
    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",

// RUN: tf-quant-opt %s -quant-insert-quantized-functions | FileCheck %s
// RUN: tf-quant-opt %s -quant-insert-quantized-functions='quantization-method=ptq target-opset=UNIFORM_QUANTIZED' --mlir-print-ir-after-all | FileCheck --check-prefix=UQ-CHECK %s

// Empty module
module {
  func.func @simple_fn(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    func.return %arg0 : tensor<*xf32>
  }
}

// CHECK-NOT: func private @internal_rescale_fn

namespace mlir::quant::stablehlo {

// Performs int8 weight-only quantization on dot_general ops.
//
// The resulting `ModuleOp` contains quantized StableHLO ops serialized in
// `TF::XlaCallModuleOp`s. They are quantized using the weight constants, not
// relying on calibration.
class WeightOnlyPtqComponent : public Component {
 public:
  // Used for debugging purposes.

  // unquantized tensors are only inserted in the unquantized model
  // whereas `DumpTensor` ops for the quantized tensors are only inserted
  // in the quantized model. Both models are required to be able to dump
  // both quantized and unquantized tensors and compare them offline.
  if (quantization_options.has_debugger_config() &&
      quantization_options.debugger_config().debugger_type() ==

  if (method.ok() && method->has_static_range_ptq()) {
    // TODO: b/331145946 - Use `Method` accessors.
    const StaticRangePtq& static_range_ptq_spec = method->static_range_ptq();
    // Look for quantized dimension specs for each quantized type and
    // populate `coeff_op_quant_dim`.
    for (const auto& [operand_idx, quantized_type] :
         static_range_ptq_spec.input_quantized_types()) {
      if (quantized_type.has_dimension_specs()) {

      std::min(quantized, static_cast<int>(std::numeric_limits<int8_t>::max()));
  quantized =
      std::max(quantized, static_cast<int>(std::numeric_limits<int8_t>::min()));
  return builder.getI32IntegerAttr(quantized);
}

// Decompose the TF ops with the registered composition library.
class DecomposeTFOpsPass
    : public PassWrapper<DecomposeTFOpsPass, OperationPass<func::FuncOp>> {
 public:

    // asymmetric range. For a state tensor, assigning correct quantization
    // parameters is sufficient, and for constants with asymmetric range it's
    // not correctly quantized by legacy quantizer so call the new Quantize.
    return Quantize(real_value, tensor_type);
  } else if (width == 16) {
    if (const auto uniform_type = dyn_cast<UniformQuantizedType>(q_type)) {
      const auto quantized_values =

std::unique_ptr<Pass> createOptimizePass();

// Creates a pass that finds quantization patterns and compose them to uniform
// quantized types.
std::unique_ptr<OperationPass<ModuleOp>>
CreateComposeUniformQuantizedTypePass();

// Creates a pass that finds stablehlo ops that accept or produce uniform
// quantized typed tensors and converts them to equivalent ops in the TFLite
// dialect.
std::unique_ptr<OperationPass<func::FuncOp>>

      }

      op.replaceAllUsesWith(q.getArg());
      return success();
    }
    return failure();
  }
};

// The StorageCastOp is used to cast from a quantized type to its storage type
// or the opposite. If none of its input and output is quantized, the op has
// no effect and should be removed.
class RemoveRedundantScast
    : public mlir::OpRewritePattern<quantfork::StorageCastOp> {
 public: