// 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]])

         is_training = false
       } : (tensor<?x64x112x112xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>)
        -> (tensor<?x64x112x112xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<*xf32>)

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: data_format = "NHWC"

  %7 = "tf.Relu"(%6) : (tensor<?x64x112x112xf32>) -> tensor<?x64x112x112xf32>
  %8 = "tf.MaxPool"(%7)
       {