}

TEST_F(IsOpFullyQuantizedTest, FalseIfOpNotQuantized) {
  constexpr absl::string_view kNotQuantizedAdd = R"mlir(
    func.func @not_quantized_add(%arg0: tensor<2xf32>) -> tensor<2xf32> {
      %0 = stablehlo.add %arg0, %arg0 : tensor<2xf32>
      return %0 : tensor<2xf32>
    }
  )mlir";

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

    quantization_max_val = 127 : i64
  } : (tensor<10xf32>, tensor<f32>, tensor<i32>) -> tensor<10x!tf_type.qint8>
  %1 = "tf.UniformDequantize"(%0, %scale, %zp) {
    quantization_axis = -1 : i64,
    quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<10x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<10xf32>
  return %1 : tensor<10xf32>
})mlir";

  %2 = func.call @fold_all_test(%arg0, %0, %1) : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
  func.return %2 : tensor<256x30x30x16xf32>
}

// ALL-LABEL: @fold_all_test
func.func @fold_all_test(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {

module {
  func.func @wrap_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
    %0 = "tf.PartitionedCall"(%arg0, %arg1) <{f = @composite_conv2d_with_relu6_fn}> {_tfl_quant_trait = "fully_quantizable"}
          : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  func.func @no_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {

  func.return %0 : tensor<1xi32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1xi32> {tf_saved_model.index_path = ["input"]}) -> (tensor<1xi32> {tf_saved_model.index_path = ["r"]})
  attributes {tf_saved_model.exported_names = ["main"]} {

      return %2 : tensor<?x2xf32>
    }
    func.func private @composite_dot_general_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {

  // CHECK-NEXT: tf_mlrt.executeop([[y]]
  %v = "tf.AddV2"(%y, %z) {__op_key = 3: i32}: (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
  return %w, %u, %v : tensor<1xi32>, tensor<1xi32>, tensor<1xi32>
}

// -----

// Test node names are preserved.

// CHECK-LABEL: func @main
func.func @main(%x: tensor<i32>) -> tensor<i32> {
  // CHECK: tf_mlrt.executeop

    %1 = "tf.BiasAdd"(%0, %cst_0) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<3xf32>) -> tensor<*xf32>
    %2 = "tf.PartitionedCall"(%arg0, %cst_2) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_1} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<*xf32>
    func.return %1, %2: tensor<*xf32>, tensor<*xf32>

func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
  // CHECK: return %arg0
  func.return %0 : tensor<?x?xf32>
}

// -----

// CHECK-LABEL: @squeeze_folder
func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<*xf32>
  // CHECK: "tfl.squeeze"

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

func.func @f32() -> tensor<4xf32> {
  // CHECK-LABEL: @f32
  // CHECK: value = dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : tensor<4xf32>
  %0 = "tfl.pseudo_const"() { value = dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32> } : () -> tensor<4xf32>
  func.return %0 : tensor<4xf32>
}

func.func @f64() -> tensor<4xf64> {
  // CHECK-LABEL: @f64