// CHECK: %[[Q:.+]] = "quantfork.qcast"(%[[CST]]) : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>
// CHECK: %[[DQ:.+]] = "quantfork.dcast"(%[[Q]]) : (tensor<2x3x3x2x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<2x3x3x2xf32>
// CHECK: %[[CALL:.+]] = "tf.XlaCallModule"(%[[ARG_0]], %[[DQ]])

  %w = arith.constant dense<2.0> : tensor<3x3x3x3xf32>
  %q = "tfl.quantize"(%w) {qtype = tensor<3x3x3x3x!quant.uniform<i8:f32, 0.1:1>>} : (tensor<3x3x3x3xf32>) -> tensor<3x3x3x3x!quant.uniform<i8:f32, 0.1:1>>
  %dq = "tfl.dequantize"(%q) : (tensor<3x3x3x3x!quant.uniform<i8:f32, 0.1:1>>) -> tensor<3x3x3x3xf32>

  %1 = "tfl.quantize"(%0) {qtype = tensor<8x!quant.uniform<u4:f32, 1.000000e+00>>} : (tensor<8xf32>) -> tensor<8x!quant.uniform<u4:f32, 1.000000e+00>>
  func.return %1 : tensor<8x!quant.uniform<u4:f32, 1.000000e+00>>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 3 : i64}>
// CHECK:  %1 = "tfl.quantize"(%0) <{qtype = tensor<8x!quant.uniform<u4:f32, 1.000000e+00>>}>
// CHECK:  return %1
}

  func.return %0, %1 : tensor<3xf32>, tensor<4x3xf32>
}

// CHECK: %[[q1:.*]] = "quantfork.qcast"(%arg0)
// CHECK-SAME: tensor<3x!quant.uniform<i8:f32, 0.0076314610593459188:-3>>
// CHECK: %[[dq1:.*]] = "quantfork.dcast"(%[[q1]])
// CHECK: %[[q2:.*]] = "quantfork.qcast"(%arg1)
// CHECK-SAME: tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {0.003937007874015748,0.0039370079913477263:-25,0.003937007874015748:51}>>
// CHECK: %[[dq2:.*]] = "quantfork.dcast"(%[[q2]])

  %1 = "quantfork.qcast"(%0) : (tensor<8xf32>) -> tensor<8x!quant.uniform<i4:f32, 1.000000e+00:-8>>
  func.return %1 : tensor<8x!quant.uniform<i4:f32, 1.000000e+00:-8>>

// CHECK:  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %cst, %cst_0) <{narrow_range = false, num_bits = 3 : i64}>
// CHECK:  %1 = "quantfork.qcast"(%0) : (tensor<8xf32>) -> tensor<8x!quant.uniform<i4:f32, 1.000000e+00:-8>>
// CHECK:  return %1
}

func.func @illegal_mhlo_uniform_quantize(%arg0: tensor<?x?xf32>) -> tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>> {
  // expected-error@+1 {{'mhlo.uniform_quantize' op is illegal as it is a UQ op or contains uq/qint types}}
  %0 = mhlo.uniform_quantize %arg0 : (tensor<?x?xf32>) -> tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>
  return %0 : tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>
}

// -----

// CHECK-NO-UNPACK: %[[QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x1024xf32>) -> tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-NO-UNPACK: %[[DOT:.+]] = stablehlo.dot_general %[[QUANTIZE_0]], %[[CONST]]

// CHECK-DAG: %[[CONST:.*]] = arith.constant dense<0.000000e+00> : tensor<2x1024xf32>
// CHECK: %0 = "quantfork.qcast"(%[[CONST]]) : (tensor<2x1024xf32>) -> tensor<2x1024x!quant.uniform<i8<-127:127>:f32, 3.9370078740157481E-9>>
// CHECK: %1 = "quantfork.dcast"(%0) : (tensor<2x1024x!quant.uniform<i8<-127:127>:f32, 3.9370078740157481E-9>>) -> tensor<2x1024xf32>

// ```
//
// The rewritten pattern looks like:
//
// ```
// %1 = stablehlo.uniform_quantize %0  // Input f32 -> uniform quantized type.
// %2 = stablehlo.constant  // Filter uniform quantized type.
// %3 = stablehlo.convolution(%1, %2)   // In uniform quantized type.
// %4 = stablehlo.uniform_dequantize %3  // Dequantize the output.
// ```
class ComposeUniformQuantizedConvolutionOp

// CHECK: %[[Q_WEIGHT_PER_TENSOR:.*]] = "quantfork.qcast"(%[[WEIGHT_CONST]]) {{.*}} (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[DQ_WEIGHT:.*]] = "quantfork.dcast"(%[[Q_WEIGHT_PER_TENSOR]])

// CHECK: %[[Q_ACTIVATION:.*]] = "quantfork.qcast"(%[[ARG_0]])
// CHECK-SAME: -> tensor<1x3x2x3x!quant.uniform<i8:f32, 0.018920717052384919:-128>>
// CHECK: %[[DQ_ACTIVATION:.*]] = "quantfork.dcast"(%[[Q_ACTIVATION]])