code: def official(x: torch.Tensor):
      # [B, C, H, W] -> [B * C * n_h, p_h, n_w, p_w]
      x = x.reshape(
          batch_size * in_channels * num_patch_h,
          patch_h,
          num_patch_w,
          patch_w,
      )
      # [B * C * n_h, p_h, n_w, p_w] -> [B * C * n_h, n_w, p_h, p_w]
      x = x.transpose(1, 2)
      # [B * C * n_h, n_w, p_h, p_w] -> [B, C, N, P] where P = p_h * p_w and N = n_h * n_w
      x = x.reshape(batch_size, in_channels, num_patches, patch_area)
      # [B, C, N, P] -> [B, P, N, C]
      x = x.transpose(1, 3)
      # [B, P, N, C] -> [BP, N, C]
      x = x.reshape(batch_size * patch_area, num_patches, -1)
  
      return x
  
label_desc: Use image node_id 1 for calling a global function with example usage: official(x) and returns: x