llvm::cl::desc("Whether enable per-channel quantized weights.")};
};

// If the weight is applicable to dynamic range quantization, insert Quantize
// and Dequantize ops with per-tensor scale.
class PrepareDRQQuantizableOp : public OpRewritePattern<arith::ConstantOp> {
 public:
  explicit PrepareDRQQuantizableOp(MLIRContext* context,

^bb0(%arg0: tensor<1x2xf32>):
  %cst_0 = arith.constant dense<[1, 0]> : tensor<2xi32>
  %0 = "tfl.quantize"(%arg0){qtype = tensor<1x2x!quant.uniform<u8:f32, 1.0>>}: (tensor<1x2xf32>) -> (tensor<1x2x!quant.uniform<u8:f32, 1.0>>)
  %1 = "tfl.dequantize"(%0): (tensor<1x2x!quant.uniform<u8:f32, 1.0>>) -> (tensor<1x2xf32>)
  %2 = "tf.Transpose"(%1, %cst_0): (tensor<1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>

  // If `value` is produced by tf.Dequantize op, then return the Dequantize op's
  // input. Otherwise return `value`.
  auto get_real_input_value = [](Value value) -> Value {
    Operation* defining_op = value.getDefiningOp();
    if (auto dequantize = dyn_cast_or_null<TF::DequantizeOp>(defining_op)) {
      return dequantize.getInput();
    } else if (auto dequantize =

      if (!read_variable_op) continue;
      // Add dequantize.
      builder.setInsertionPointAfter(read_variable_op);
      auto new_read_variable_op =
          builder.create<ReadVariableOp>(read_variable_op.getLoc(), ref_qtype,
                                         read_variable_op.getResourceId());
      auto new_dq_op = builder.create<DequantizeOp>(

        per_axis_type.getStorageTypeMin(), per_axis_type.getStorageTypeMax());
  }

  auto quantize = builder.create<quantfork::QuantizeCastOp>(
      q_op.getLoc(), new_value_type.clone(new_qtype), new_value);
  auto dequantize = builder.create<quantfork::DequantizeCastOp>(
      dq_op.getLoc(), new_value_type, quantize.getResult());
  return dequantize.getResult();
}

    * A tensor is dequantized using a `func::FuncOp` whose name contains
      "uniform_dequantize". The first argument is the tensor to be quantized,
      the second argument is the zero point constant (element type: int) and
      the third argument is the inverse scale constant (element type: float).
    * Inputs to the target quantized op is quantized and the outputs are
      dequantized.

  ];
}
def DecomposeHybridQuantizationPass : Pass<"tfl-decompose-hybrid-quantization", "mlir::func::FuncOp"> {
  let summary = "Decomposes hybridge quantization to explicit quantize / dequantize";
  let description = [{
      Decomposes (with explicit quantize/dequantize ops) selected math
      operations which exist in the model with hybrid quantization
      (some arguments/results left in floating point).
  }];

    op->erase();
  });
}

// Fold quantized i32 (normally bias) into their float values.
struct FoldQuantizedI32ToFloat : public OpRewritePattern<TFL::DequantizeOp> {
  using OpRewritePattern<TFL::DequantizeOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::DequantizeOp dequant_op,
                                PatternRewriter& rewriter) const override {
    // We only fold i32 -> float pattern.

// 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
// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type);

// Returns true if a type is quantized tensor type.
bool IsQuantizedTensorType(Type type);

}
// CHECK-LABEL: uniform_dequantize_op_ui16_storage_input
// CHECK: stablehlo.uniform_dequantize
// CHECK-NOT: tfl.dequantize

// -----

// Tests that the pattern doesn't match when the input quantized tensor's
// storage type is i32. i32 storage type is not compatible with
// `tfl.dequantize`.

func.func @uniform_dequantize_op_i32_storage_input(%arg: tensor<2x2x!quant.uniform<i32:f32, 1.000000e+0:8>>) -> tensor<2x2xf32> {