ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op) {
  auto composite_result_shape =
      mlir::cast<ShapedType>(old_op->getResults().front().getType()).getShape();
  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW
  output_shape[0] = composite_result_shape[0];
  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];
  output_shape[3] = composite_result_shape[1];