// CHECK:               : (tensor<1x114x114x64xf32>) -> tensor<1x64x114x114xf32>
// CHECK:           return %[[TPOS]] : tensor<1x64x114x114xf32>

func.func @commute_transpose_pad(
      %arg0: tensor<1x112x112x64xf32>, %padding_val: tensor<f32>)
      -> tensor<1x64x114x114xf32> {
  %tspos = stablehlo.transpose %arg0, dims = [0, 3, 1, 2]
      : (tensor<1x112x112x64xf32>) -> tensor<1x64x112x112xf32>
  %ret = stablehlo.pad %tspos, %padding_val,

  // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[MAX_POOL]], %[[RES_PERM]])
  // CHECK: return %[[RES_TRANSPOSE]]

  // Transpose NCHW -> NHWC
  %0 = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi32>} : () -> tensor<4xi32>
  %1 = "tf.Transpose"(%arg0, %0) : (tensor<1x64x112x112xf32>, tensor<4xi32>) -> tensor<1x112x112x64xf32>

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

// CHECK-LABEL: maxpool_nchw
func.func @maxpool_nchw(%arg0: tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32> {
  // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[R0:.*]] = "tf.Transpose"(%arg0, %[[CST]])

  %1 = "tf.Conv2D"(%input, %filter)
       {
         data_format = "NCHW",
         padding = "VALID",
         strides = [1, 1, 2, 2]
       } : (tensor<1x64x28x28xf32>, tensor<1x1x64x64xf32>)
        -> tensor<1x64x14x14xf32>

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

// CHECK-LABEL: func @transposeConv2D_3x3_f16