}

// -----

// CHECK-LABEL: func @convolution_add_add_i8
func.func @convolution_add_add_i8(
    %lhs: tensor<?x3x2x1xi8>, %rhs: tensor<2x1x1x1xi8>,
    %zp_offset: tensor<?x2x2x1xi8>, %bias: tensor<1xi8>
  ) -> tensor<?x2x2x1xi8> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]

    %11 = stablehlo.reshape %2 : (tensor<1x1x1x1xi8>) -> tensor<1xi8>
    %12 = stablehlo.broadcast_in_dim %11, dims = [0] : (tensor<1xi8>) -> tensor<1x3x3x4xi8>
    %13 = stablehlo.convert %12 : (tensor<1x3x3x4xi8>) -> tensor<1x3x3x4xf32>
    %14 = stablehlo.convert %3 : (tensor<3x3x4x4xi8>) -> tensor<3x3x4x4xf32>

  func.return %0 : tensor<4xf64>
}

func.func @i8() -> tensor<4xi8> {
  // CHECK-LABEL: @i8
  // CHECK: value = dense<[1, 2, 3, 4]> : tensor<4xi8>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 4]> : tensor<4xi8> } : () -> tensor<4xi8>
  func.return %0 : tensor<4xi8>
}

func.func @i16() -> tensor<4xi16> {
  // CHECK-LABEL: @i16

  func.return %0 : tensor<4xf64>
}

func.func @i8() -> tensor<4xi8> {
  // CHECK-LABEL: @i8
  // CHECK: value = dense<[1, 2, 3, 4]> : tensor<4xi8>
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 4]> : tensor<4xi8> } : () -> tensor<4xi8>
  func.return %0 : tensor<4xi8>
}

func.func @i16() -> tensor<4xi16> {
  // CHECK-LABEL: @i16

  // CHECK: %[[CONVERT_1:.*]] = mhlo.bitcast_convert %[[QUANTIZE]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xi8>
  // CHECK: %[[CONVERT_2:.*]] = mhlo.bitcast_convert %[[CONVERT_1]] : (tensor<2xi8>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>
  // CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_2]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xf32>

}

// CHECK-LABEL: @cast_ui8_to_i8
func.func @cast_ui8_to_i8() -> tensor<4xi8> {
  %cst = arith.constant dense<[0, 255, 127, 128]> : tensor<4xui8>
  %0 = "tfl.cast"(%cst) : (tensor<4xui8>) -> tensor<4xi8>
  func.return %0 : tensor<4xi8>

// CHECK: %[[CST:.*]] = arith.constant dense<[0, -1, 127, -128]> : tensor<4xi8>
// CHECK:  return %[[CST]]
}

// CHECK-LABEL: @cast_i8_to_i32

  func.return %11 : tensor<*xf32>
}

// CHECK: %[[maxpool_i8:.*]] = "tf.MaxPool"
// CHECK-SAME: (tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xi8>, tensor<2xi32>) -> tensor<*xi8>
// CHECK: %[[scast:.*]] = "quantfork.scast"(%[[reshape_i8]]
// CHECK: %[[matmul:.*]] = "tf.PartitionedCall"(%[[scast]]
// CHECK-SAME: f = @composite_matmul_fn_1

  // Argument has rank 2, but its dimensions are unknown.
  constexpr absl::string_view kArgumentWithUnknownDims = R"mlir(
    func.func @unknown_dims_arg(%arg: tensor<?x?xi8>) -> tensor<?x?xi8> {
      return %arg : tensor<?x?xi8>
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kArgumentWithUnknownDims);
  ASSERT_TRUE(module_op);

// test stablehlo roundtrip

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

func.func @logistic(%arg0: tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32> {

}

TEST_F(ConvertTfQuantToMhloIntTest, UniformRequantizeToValidGraph) {
  constexpr absl::string_view kProgram = R"mlir(
func.func @main(%input: tensor<10xi8>) -> tensor<10xi8> {
  %input_scale = "tf.Const"() { value = dense<0.2235> : tensor<f32> } : () -> tensor<f32>
  %input_zp = "tf.Const"() { value = dense<-2> : tensor<i32> } : () -> tensor<i32>