// CHECK:              : (tensor<1x56x56x64xf32>) -> tensor<1x64x56x56xf32>
// CHECK:           return %[[TPOS]] : tensor<1x64x56x56xf32>

func.func @commute_transpose_reduce_window(
      %input: tensor<1x114x114x64xf32>,
      %cst: tensor<f32>) -> tensor<1x64x56x56xf32> {
  %tpos = stablehlo.transpose %input, dims = [0, 3, 1, 2]
      : (tensor<1x114x114x64xf32>) -> tensor<1x64x114x114xf32>

       {
         data_format = "NCHW",
         ksize = [1, 1, 3, 3],
         padding = "SAME",
         strides = [1, 1, 2, 2]
       } : (tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32>
  func.return %0 : tensor<1x64x56x56xf32>

  // CHECK: %[[MAX_POOL:[0-9]*]] = "tf.MaxPool"(%arg0) <{data_format = "NCHW", ksize = [1, 1, 3, 3], padding = "SAME", strides = [1, 1, 2, 2]}> : (tensor<1x64x112x112xf32>) -> tensor<1x64x56x56xf32>
  // CHECK: %[[RES_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%[[MAX_POOL]], %[[RES_PERM]])
  // CHECK: return %[[RES_TRANSPOSE]]

  // Transpose NCHW -> NHWC

  %8 = "tf.MaxPool"(%7)
       {
         data_format = "NCHW",
         ksize = [1, 1, 3, 3],
         padding = "SAME",
         strides = [1, 1, 2, 2]
       } : (tensor<?x64x112x112xf32>) -> tensor<?x64x56x56xf32>

  // CHECK: %[[MAX_POOL:[0-9]*]] = "tf.MaxPool"
  // CHECK-SAME: data_format = "NHWC"
  // CHECK-SAME: ksize = [1, 3, 3, 1]
  // CHECK-SAME: strides = [1, 2, 2, 1]

func.func @upsample_bilinear2d(%arg0: tensor<1x64x16x16xf32>) -> (tensor<1x64x32x32xf32>) {
  %0 = mhlo.composite "odml.upsample_bilinear2d" %arg0 {composite_attributes = {is_nchw_op = true, align_corners = false, output = dense<32> : tensor<2xi64>}, decomposition = @XlaCallModule_odml.upsample_bilinear2d.impl_21_0} : (tensor<1x64x16x16xf32>) -> tensor<1x64x32x32xf32>
  return %0 : tensor<1x64x32x32xf32>
}