%q1 = "quantfork.qcast"(%arg0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 6.000000e-03:-128>>
  %q2 = "quantfork.qcast"(%cst) {volatile} : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i8<-127:127>:f32, 5.000000e-03>>

  %cst = "tfl.pseudo_qconst"() {qtype = tensor<1x2x!quant.uniform<u8:f32, 1.0>>, value = dense<-76> : tensor<1x2xi8>} : () -> tensor<1x2x!quant.uniform<u8:f32, 1.0>>

// CHECK: %[[quant:.*]] = "tfr.cast"(%arg0) : (tensor<1x10xi8>) -> !tfr.tensor
// CHECK: %[[scale_cast:.*]] = "tfr.cast"(%[[scale]])
// CHECK: %[[zp_cast:.*]] = "tfr.cast"(%[[zp]])
// CHECK: %[[requant:.*]] = tfr.call @tf__requantize(%[[quant]], %[[scale_cast]], %[[zp_cast]])
// CHECK: %[[result:.*]] = "tfr.cast"(%[[requant]])
// CHECK-NOT: quant.uniform
// CHECK: return %[[result]], %[[float_resunt:.*]] : tensor<1x10xi8>, tensor<1x5xf32>
}

    return %0 : tensor<1x3xf32>
  }
}

// CHECK-LABEL: quantize_dot_general_per_tensor
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>
// CHECK: %[[CST:.+]] = stablehlo.constant() <{value = dense<127> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>

// Check that the `attr_map` attribute has been removed.
// CHECK-NOT: attr_map

// -----

func.func @conv_with_dynamic_channel_dim(%arg0: tensor<1x3x4x?xf32>) -> tensor<*xf32> {
  %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<1xf32>} : () -> tensor<1xf32>

  %1 = "quantfork.qcast"(%0) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>
  %2 = "quantfork.dcast"(%1) : (tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>) -> tensor<1x3xf32>
  return %2 : tensor<1x3xf32>
}
// CHECK: %[[XLA_CALL_MODULE_0:.+]] = "tf.XlaCallModule"() <{{{.*}}}> {{{.*}}} : () -> tensor<1x3x!quant.uniform<i8:f32, 1.000000e-03:-3>>

// RUN: stablehlo-quant-opt %s -verify-quant-legalization -split-input-file -verify-diagnostics | FileCheck %s

// CHECK-LABEL: func @legal_i8
func.func @legal_i8(%arg0: tensor<1xi8>) -> tensor<1xi8> {
  func.return %arg0: tensor<1xi8>
}

// -----

func.func @illegal_qint8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1x!tf_type.qint8> {
  // expected-error@+1 {{'func.return' op is illegal as it is a UQ op or contains uq/qint types}}

// CHECK-NO-UNPACK: %[[QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[DOT]] : (tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-NO-UNPACK: %[[DEQUANTIZE:.+]] = stablehlo.uniform_dequantize %[[QUANTIZE_1]] : (tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>) -> tensor<1x3xf32>
// CHECK-NO-UNPACK: return %[[DEQUANTIZE]] : tensor<1x3xf32>

// -----

  %weight_zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

  // CHECK: "tf.AddV2"
  // CHECK: mhlo.constant
  // CHECK-SAME{LITERAL}: dense<[[1, 2], [3, 4]]> : tensor<2x2xi8>
  %0 = "tf.AddV2"(%input, %input) : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>

  // CHECK: "mhlo.dot"
  // CHECK-SAME: (tensor<?x?xf32>, tensor<2x2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?x?xf32>