func.func @cast_op_qint8_int8(%arg0: tensor<1x!tf_type.qint8>) -> tensor<1xi8> {
  // CHECK: "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<1xi8>) -> tensor<1xi8>
  %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
  func.return %0: tensor<1xi8>
}

// -----

// CHECK-LABEL: func @cast_op_int8_qint8
func.func @cast_op_int8_qint8(%arg0: tensor<1xi8>) -> tensor<1x!tf_type.qint8> {

}

// CHECK-LABEL: @cast_i32_to_ui8
func.func @cast_i32_to_ui8() -> tensor<6xui8> {
  %cst = arith.constant dense<[0, -1, 256, 127, -128, -129]> : tensor<6xi32>
  %0 = "tfl.cast"(%cst) : (tensor<6xi32>) -> tensor<6xui8>
  func.return %0 : tensor<6xui8>

// CHECK: %[[CST:.*]] = arith.constant dense<[0, 255, 0, 127, 128, 127]> : tensor<6xui8>
// CHECK:  return %[[CST]]
}

// CHECK-LABEL: @cast_ui8_to_i8

// 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}}

  func.return %0 : tensor<!tf_type.string>
}

func.func @uint8() -> tensor<4xui8> {
  // CHECK-LABEL: @uint8
  // CHECK: value = dense<[222, 173, 190, 239]> : tensor<4xui8>
  %0 = "tfl.pseudo_const"() {value = dense<[222, 173, 190, 239]> : tensor<4xui8>} : () -> tensor<4xui8>
  func.return %0 : tensor<4xui8>
}

func.func @int4() -> tensor<5xi4> {
  // CHECK-LABEL: @int4

  producer: 29
  min_consumer: 12
}

# CHECK: func @main
# CHECK: "tf.PartitionedCall"()
# CHECK-SAME: f = @[[FUNCTION:[A-Za-z0-9_]*]]
# CHECK: func private @[[FUNCTION]]() -> tensor<*xui8>

func.func @shift_right_unsigned(%arg0: tensor<4xui8>, %arg1: tensor<4xui8>) -> tensor<4xui8> {
  // CHECK:  mhlo.shift_right_logical %arg0, %arg1 : tensor<4xui8>
  %0 = "tf.RightShift"(%arg0, %arg1) : (tensor<4xui8>, tensor<4xui8>) -> tensor<4xui8>
  func.return %0 : tensor<4xui8>
}

// CHECK-LABEL: func @broadcast_shift_right_unsigned
func.func @broadcast_shift_right_unsigned(%arg0: tensor<4xui8>, %arg1: tensor<2x4xui8>) -> tensor<2x4xui8> {

  func.return %0 : tensor<!tf_type.string>
}

func.func @uint8() -> tensor<4xui8> {
  // CHECK-LABEL: @uint8
  // CHECK: value = dense<[222, 173, 190, 239]> : tensor<4xui8>
  %0 = "tfl.pseudo_const"() {value = dense<[222, 173, 190, 239]> : tensor<4xui8>} : () -> tensor<4xui8>
  func.return %0 : tensor<4xui8>
}

func.func @int4() -> tensor<5xi4> {
  // CHECK-LABEL: @int4

// CHECK: %[[BITCAST_CONVERT_0:.*]] = stablehlo.bitcast_convert %[[UQ]] : (tensor<1x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<1xi8>
// CHECK: %[[BITCAST_CONVERT_1:.*]] = stablehlo.bitcast_convert %[[BITCAST_CONVERT_0]] : (tensor<1xi8>) -> tensor<1x!quant.uniform<i8:f32, 1.000000e+00:3>>

  %outputs = "tf.Const"() {value = dense<224> : tensor<2xi32>} : () -> tensor<2xi32>

// CHECK:       %[[POP_CNT:.*]] = "tf.PopulationCount"(%[[ARG_0]]) : (tensor<8xui8>) -> tensor<8xui8>
// CHECK:       return %[[POP_CNT]]
// CHECK:         }
func.func @convert_population_count_ui8(%arg0: tensor<8xui8>) -> tensor<8xui8> {
  %0 = "mhlo.popcnt"(%arg0) : (tensor<8xui8>) -> tensor<8xui8>
  func.return %0 : tensor<8xui8>
}