// RUN: odml-to-stablehlo-opt --compose-uniform-quantized-type \
// RUN:     --split-input-file --verify-diagnostics %s | FileCheck %s

module {
// CHECK-LABEL: quantized_conv_op
// CHECK-SAME: %[[ARG:.*]]: tensor<1x3x3x4xf32>
  func.func @quantized_conv_op(%arg0: tensor<1x3x3x4xf32>) -> tensor<1x3x3x4xf32> {
    %1 = stablehlo.constant dense<1.000000e+03> : tensor<1x1x1x1xf32>  // Input inverse scale.

// RUN: odml-to-stablehlo-opt --uniform-quantized-stablehlo-to-tfl \
// RUN:     --split-input-file --verify-diagnostics %s | FileCheck %s

// ============================================================================
// The following functions tests example quantization patterns outputted from
// JAX Quantizer. JAX Quantizer should output integer types, which are
// composed into `UniformQuantized{|PerAxis}Type` via
// `compose_uniform_quantized_type_pass.cc`.

          !IsConnectedWithCompsiteFunction(quantizing_op)) {
        continue;
      }

      // Same scale op is not supported for Uniform Quantized ops.
      if (target_opset_ == OpSet::UNIFORM_QUANTIZED) {
        continue;
      }

      // Collect all the quantized inputs and "clone" the matched op by these
      // inputs.
      SmallVector<Value, 4> inputs;
      inputs.reserve(quantizing_op->getNumOperands());

// If the selected quantization option is not available, StableHLO quantizer
// will raise an error.
// NEXT ID: 2
message CustomQuantizationMethod {
  // Specify component name, bit width, and other specs for all compoenents
  // intended to be quantized.
  repeated QuantizationComponentSpec quantization_component_spec = 1;
}

// Quantization spec per each component designated to be quantized.

  string name = n;
  string asTraitArgsStr =
    !interleave(params, ", ") # !if(signed, ", true", ", false");
}

// Uniform quantized types. Two integers "smantissa" and "sexp" are used to
// express the Mantissa and Exponent components of the floating-point scale so
// the scale of the quantized type is "smantissa * 10 ^ sexp".
class UInt8UniformQuantizedType<int zero_pt, int smantissa, int sexp>
    : QuantizedType<"Uniform",

// Returns true iff `type` is a uniform quantized type whose storage type is
// 32-bit integer and expressed type is f32.
bool IsI32F32UniformQuantizedType(Type type);

// Returns true iff `type` is a uniform quantized per-axis (per-channel) type
// whose storage type is 32-bit integer and expressed type is f32.
bool IsI32F32UniformQuantizedPerAxisType(Type type);

// Determines whether the storage type of a quantized type is supported by

}

def QuantizeCompositeFunctionsPass : Pass<"stablehlo-quantize-composite-functions", "ModuleOp"> {
  let summary = "Quantize composite functions with QDQ input / outputs.";
  let options = [
    Option<"enable_per_channel_quantized_weight_",
        "enable-per-channel-quantized-weight",
        "bool", /*default=*/"true",
        "Whether to enable per-channel quantized weights.">,
    Option<"mlir_dump_file_name_", "mlir-dump-file-name",

//
//   1. Replaces quantized `TF::XlaCallModuleOp` with a `func::CallOp`.
//   2. Quantizes the callee function.
//
// The inputs of this pattern assumes an invalid IR, where even if a
// `TF::XlaCallModuleOp` is quantized the callee remains unquantized. Step (2)
// not only replaces the input and output tensor types into quantized ones, but
// also rewrites the body with a quantized equivalent.
//

};

// Converts quantized tensor type with signed integer type to quantized tensor
// type with unsigned integer type.
Type ConvertSignedQuantizedToUnsigned(Type signed_tensor_type, Location loc);

// Converts quantize ops with unsigned quantized types to these with signed
// quantized types and preserves the scales.
template <typename QuantizeOpT>

// CHECK: -------- Quantization Summary --------
// CHECK: Number of quantized layers in the model
// CHECK: --------------------------------
// CHECK: Name    Count/Total
// CHECK: ================================
// CHECK: Conv2D  1/1

// CHECK: Number of quantized layers with quantized outputs: 0/1
// CHECK: Number of quantize layers added: 1
// CHECK: Number of dequantize layers added: 0
}

// -----