%1 = "tf.Transpose"(%arg0, %0) : (tensor<1x3x32x32xf32>, tensor<4xi32>) -> tensor<1x32x32x3xf32>

  // Compute in NHWC
  %2 = "tf.Conv2D"(%1, %arg1)
    {
      data_format = "NHWC",
      padding = "SAME",
      strides = [1, 1, 1, 1],
      dilations = [1, 1, 1, 1]
    } : (tensor<1x32x32x3xf32>, tensor<1x1x3x8xf32>) -> tensor<1x32x32x8xf32>

  // Convert result back: NHWC -> NCHW

  func.return %2 : tensor<1x32x32x16x!quant.uniform<i8:f32, 1.0>>
}

// -----

// CHECK-LABEL: @test_conv2d_replace_float
func.func @test_conv2d_replace_float(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x16xf32> {

         strides = [5, 6, 7, 8]
       } : (tensor<1x32x32x3xf32>, tensor<4xi32>, tensor<1x32x32x8xf32>)
         -> tensor<1x1x3x8xf32>

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

// CHECK-LABEL: func @transposeConv2DBackpropInput
func.func @transposeConv2DBackpropInput(
  %input_sizes: tensor<4xi32>,
  %filter: tensor<1x1x3x8xf32>,
  %out_backprop: tensor<1x32x32x8xf32>
) -> tensor<1x32x32x3xf32> {

}

// -----

func.func @testBiasAddNoDataFormatOk(tensor<1x32x32x16xf32>, tensor<16xf32>) -> tensor<1x32x32x16xf32> {
^bb0(%arg0: tensor<1x32x32x16xf32>, %arg1: tensor<16xf32>):
  %0 = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT"}: (tensor<1x32x32x16xf32>, tensor<16xf32>) -> tensor<1x32x32x16xf32>
  func.return %0 : tensor<1x32x32x16xf32>
}

// -----

      // CHECK-OPT: %{{.*}} = stablehlo.custom_call @tf.ResizeBilinear(%arg0, %cst) {align_corners = false, device = "", half_pixel_centers = true} : (tensor<1x32x32x128xf32>, tensor<2xi32>) -> tensor<1x64x64x128xf32>
      %1 = "tf.ResizeBilinear"(%arg0, %0) {

  // CHECK: %[[VAL_1:.*]] = "tf.Identity"(%[[VAL_0]]) : (tensor<*xf32>) -> tensor<*xf32>
  // CHECK: return %[[VAL_1]]
  %0 = "tf.Conv2D"(%arg2, %arg1) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>) -> tensor<*xf32>

  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> tensor<3x32x32x3xf32> {
      %0 = mhlo.constant dense<2.550000e+02> : tensor<3x32x32x3xf32>
      return %0 : tensor<3x32x32x3xf32>
    }
  })";
  CreateModule(kMlirModuleStr);

  auto result = Run();

  EXPECT_TRUE(result.failed());
}

}  // namespace

// CHECK:  %1 = "tfl.transpose"(%arg0, %cst) : (tensor<4x8x32x32xf32>, tensor<4xi32>) -> tensor<4x32x32x8xf32>
// CHECK:  %cst_0 = arith.constant dense<64> : tensor<2xi32>
// CHECK:  %2 = "tfl.resize_nearest_neighbor"(%1, %cst_0) <{align_corners = false, half_pixel_centers = true}> : (tensor<4x32x32x8xf32>, tensor<2xi32>) -> tensor<4x64x64x8xf32>
// CHECK:  %cst_1 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>

  func.return %0 : tensor<256x32x32x16xf32>
}

func.func @testConv2DDynamicShape(tensor<?x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<?x32x32x16xf32> {
^bb0(%arg0: tensor<?x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>):
  // CHECK: _arithmetic_count = -1 : i64

// RUN: tf-opt %s -canonicalize | FileCheck %s

// CHECK-LABEL: func @testShape
func.func @testShape(tensor<f32>, tensor<1x32x32x16xf32>, tensor<*xf32>) -> (tensor<0xi32>, tensor<?xi32>, tensor<?xi32>) {
^bb0(%arg0: tensor<f32>, %arg1: tensor<1x32x32x16xf32>, %arg2: tensor<*xf32>):

  // CHECK-DAG: tf.Const{{.*}} dense<> : tensor<0xi32>
  %0 = "tf.Shape"(%arg0) {T = "tfdtype$DT_FLOAT", output = "tfdtype$DT_INT32"} : (tensor<f32>) -> tensor<0xi32>