// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[CONV]], dims = [0, 3, 1, 2] : (tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32>
// CHECK: return %[[TRANSPOSE_1]]

// -----

// Tests that a `add(convolution(%activation, %weight), %bias)` with the
// activation tensor of NCHW format is converted to NHWC convolution + add

  %result = "tf_device.launch"() ({
    %3 = "tf.Transpose"(%2, %1) : (tensor<1x8x4x4xf32>, tensor<4xi32>) -> tensor<1x4x4x8xf32>
    tf_device.return %3: tensor<1x4x4x8xf32>
  }) {device = "device"} : () -> tensor<1x4x4x8xf32>

  func.return %result : tensor<1x4x4x8xf32>

  // CHECK-DAG: %[[CONST1:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}>

  func.return %rst : tensor<256x30x30x16xf32>

// CHECK-DAG: %[[CONSTANT:.*]] = arith.constant dense<0.000000e+00> : tensor<48xf32>
// CHECK-DAG: %[[CONSTANT0:.*]] = arith.constant dense<0.000000e+00> : tensor<1x3x3x48xf32>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<1x3x3x48x!quant.uniform<u8:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])

    rhs_quantization_min_val = -128 : i64,
    rhs_quantization_max_val = 127 : i64
  } : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x3x3x2xf32>
  func.return %0 : tensor<1x3x3x2xf32>
}

//===----------------------------------------------------------------------===//
// tf.UniformQuantize and tf.UniformDequantize legalization

  func.return %rst : tensor<256x30x30x16xf32>

// CHECK-DAG: %[[CONSTANT:.*]] = arith.constant dense<0.000000e+00> : tensor<48xf32>
// CHECK-DAG: %[[CONSTANT0:.*]] = arith.constant dense<0.000000e+00> : tensor<1x3x3x48xf32>
// CHECK: %[[QUANTIZE:.*]] = "tfl.quantize"(%[[CONSTANT0]]) <{qtype = tensor<1x3x3x48x!quant.uniform<u4:f32, 1.000000e+00>>}>
// CHECK: %[[DEQUANTIZE:.*]] = "tfl.dequantize"(%[[QUANTIZE]])

func.func @DontFuseAddWithConvActivationFunc(%arg0: tensor<1x3x1x1xf32>) -> tensor<1x2x1x3xf32> {
  %cst = arith.constant dense<1.5> : tensor<1xf32>
  %cst_1 = arith.constant dense<0.0> : tensor<3xf32>
  %cst_2 = arith.constant dense<1.1> : tensor<3x2x1x1xf32>
  %0 = "tfl.add"(%arg0, %cst) {fused_activation_function = "RELU6"} : (tensor<1x3x1x1xf32>, tensor<1xf32>) -> tensor<1x3x1x1xf32>

// CHECK-LABEL: func @testValidFakeQuantWithMinMaxArgs
func.func @testValidFakeQuantWithMinMaxArgs(tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32> {
^bb0(%arg0: tensor<8x8x8x8xf32>):
  %0 = "tf.FakeQuantWithMinMaxArgs"(%arg0) {min = -1.0 : f32, max = 1.0 : f32, num_bits = 3} : (tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32>
  func.return %0 : tensor<8x8x8x8xf32>
}

// -----

// Test invalid tf.FakeQuantWithMinMaxArgs

}

// -----

func.func @transpose_uniform_qtype(%arg0 : tensor<1x3x4x3xf32>) -> tensor<3x4x3x1x!quant.uniform<i8:f32, 0.0078356266021728515:-1>> {
  %cst = arith.constant dense<[1, 2, 3, 0]> : tensor<4xi32>
  %0 = "tfl.quantize"(%arg0) {qtype = tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0078356266021728515:-1>>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3x!quant.uniform<i8:f32, 0.0078356266021728515:-1>>
  // CHECK: "tfl.transpose"