// CHECK: return %[[res]] : tensor<2xi1>
}

// CHECK-LABEL: fixed_element_attribute_invalid
func.func @fixed_element_attribute_invalid(%arg0: tensor<2xf32>) -> tensor<2xi8> {
  %0 = "tfr.cast"(%arg0) : (tensor<2xf32>) -> !tfr.tensor
  // expected-error@+1 {{type i8_ can't be resolved for the signature of the op}}

func.func @testI8ToI4WithinRange() -> (tensor<4xi8>) {
  %0 = arith.constant dense<[-8, 0, 1, 7]> : tensor<4xi8>
  // expected-error@+1 {{type of return operand 0 ('tensor<4xi4>') doesn't match function result type ('tensor<4xi8>')}}
  func.return %0 : tensor<4xi8>
}

func.func @testI8ToI4NotWithinRange() -> tensor<4xi8> {
  %0 = arith.constant dense<[-10, 2, 3, 8]> : tensor<4xi8>
  func.return %0 : tensor<4xi8>

// CHECK-LABEL: @main
func.func @main(%arg0: tensor<4xi8>, %arg1: tensor<4xi32>) -> (tensor<4xi32>, tensor<4xi8>)
// CHECK: attributes {tf.entry_function = {inputs = "input0,input1", outputs = "output0,output1"}}
attributes {tf.entry_function = {inputs = "input0,input1", outputs = "output0,output1"}} {
  %0 = "tfl.neg"(%arg0) : (tensor<4xi8>) -> tensor<4xi8> loc("neg")
  %1 = "tfl.neg"(%arg1) : (tensor<4xi32>) -> tensor<4xi32> loc("neg")

  %a = corert.const_dense_tensor dense<[true, false]> : tensor<2xi1>
  %b = corert.const_dense_tensor dense<[false, true]> : tensor<2xi1>
  // CHECK: [[b:%.*]] = tfrt_fallback_async.const_dense_tensor dense<[false, true]> : tensor<2xi1>
  // CHECK-NEXT:      [[a:%.*]] = tfrt_fallback_async.const_dense_tensor dense<[true, false]> : tensor<2xi1>
  // CHECK-NEXT: tfrt_fallback_async.corert_tensorhandle_to_fallback_tensor [[arg0]]

// Identity function to make the exporter happy
func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> {
  func.return %arg0 : tensor<4xi8>
}

//CHECK:func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> attributes {tf.entry_function = {inputs = "arg0", outputs = "arg0"}} {
//CHECK: return %arg0 : tensor<4xi8>
//CHECK:}

// 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.func private @internal_conv3d_fn(
                         %input : tensor<*xf32>, %filter : tensor<*xi8>) -> tensor<*xf32> {

   // Use identity op to avoid the filter being constant-folded.
    %identity = "tf.Identity"(%filter) : (tensor<*xi8>) -> tensor<*xi8>
    %2 = "tf.Cast"(%identity) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32>
    %3 = "tf.Conv3D"(%input, %2) {
      padding = "VALID", strides = [1, 1, 1, 1, 1],

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

    return %3 : tensor<1x2x2x1024xbf16>
  }

  func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xbf16> {
    %cst = "tf.Const"() {value = dense<1.574710e-02> : tensor<bf16>} : () -> tensor<bf16>
    %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<*xi8>) -> tensor<*xbf16>
    %1 = "tf.Mul"(%0, %cst) : (tensor<*xbf16>, tensor<bf16>) -> tensor<*xbf16>
    return %1 : tensor<*xbf16>
  }

// CHECK: %[[maxpool_i8:.*]] = "tf.MaxPool"(%[[sc1]])
// CHECK-SAME: (tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xi8>, tensor<2xi32>) -> tensor<*xi8>
// CHECK: %[[sc2:.*]] = "quantfork.scast"(%[[reshape_i8]])
// CHECK: %[[dq:.*]] = "quantfork.dcast"(%[[sc2]]) : (tensor<*x!quant.uniform<i8:f32, 5.000000e-02:-10>>)