// CHECK: return %[[UNIFORM_QUANTIZE_1]] : tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>

// CHECK-PER-TENSOR: func.func private @quantized_dot_general_fn(%[[ARG_1:.+]]: tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>, %[[ARG_2:.+]]: tensor<2x3x!quant.uniform<i8<-127:127>:f32, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>> attributes {_from_xla_call_module}

// 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-NEXT: %[[q2:.*]] = "tfl.quantize"(%[[cst_0]]) <{qtype = tensor<1x3x!quant.uniform<i8<-127:127>:f32:0, {{.*}}>>}> {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8<-127:127>:f32:0, {{.*}}>>
// CHECK-NEXT: %[[dq2:.*]] = "tfl.dequantize"(%[[q2]]) : (tensor<1x3x!quant.uniform<i8<-127:127>:f32:0, {{.*}}>>) -> tensor<1x3xf32>

  %2 = "tf.Fill"(%1, %cst_0) {device = ""} : (tensor<2xi32>, tensor<i32>) -> tensor<1x3xi32>

  // Just to introduce extra usages for %cst.
  %3 = "tf.Mul"(%cst, %cst) {device = ""} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>

  return %2: tensor<1x3xi32>
}
// The output of tf.Shape is considered a compile-time constant, so the

    // CHECK: %[[Q3:.*]] = "quantfork.qcast"(%[[ARG3]]) {volatile} : (tensor<1x3xf32>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
    // CHECK: %[[SELECT:.*]] = stablehlo.select %[[ARG2]], %[[CALL]], %[[Q3]] : tensor<1x3xi1>, tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>

  // CHECK-LABEL: @qu8
  // CHECK: <{qtype = tensor<3x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>, value = dense<1> : tensor<3xi8>}> : () -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>
  %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>, value = dense<1> : tensor<3xi8>} : () -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
  func.return %0 : tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
}

  // CHECK-LABEL: @qu8
  // CHECK: <{qtype = tensor<3x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>, value = dense<1> : tensor<3xi8>}> : () -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.000000e+00>>
  %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>, value = dense<1> : tensor<3xi8>} : () -> tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
  func.return %0 : tensor<3x!quant.uniform<u8<1:255>:f32, 1.0>>
}

// TF-DAG: %[[cc_0:.*]] = "tf.ConcatV2"(%[[pc_1]], %[[pc_5]], %[[cst_0]]) : (tensor<2x2xi8>, tensor<2x2xi8>, tensor<i32>) -> tensor<2x4xi8>
// TF-DAG: %[[pc_7:.*]] = "tf.PartitionedCall"(%[[cc_0]], %[[cst_6]], %[[cst_5]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}
// TF-DAG: return %[[pc_7]]