code: def window_partition(
      x: torch.Tensor, window_size: int
  ) -> Tuple[torch.Tensor, Tuple[int, int]]:
      """
      Partition into non-overlapping windows with padding if needed.
      Args:
          x (tensor): input tokens with [B, H, W, C].
          window_size (int): window size.
  
      Returns:
          windows: windows after partition with [B * num_windows, window_size, window_size, C].
          (Hp, Wp): padded height and width before partition
      """
      B, H, W, C = x.shape
  
      pad_h = (window_size - H % window_size) % window_size
      pad_w = (window_size - W % window_size) % window_size
      x = F.pad(x, (0, 0, 0, pad_w, 0, pad_h))
      Hp, Wp = H + pad_h, W + pad_w
  
      x = x.view(
          B,
          Hp // window_size,
          window_size,
          Wp // window_size,
          window_size,
          C,
      )
      windows = (
          x.permute(0, 1, 3, 2, 4, 5)
          .contiguous()
          .view(-1, window_size, window_size, C)
      )
      return windows, (Hp, Wp)
  
imports: import logging
  from functools import partial
  from typing import Callable, Optional, Tuple
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.utils.checkpoint
  from torch.jit import Final
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD, IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, PatchDropout, LayerNorm2d, ClassifierHead, NormMlpClassifierHead, Format, resample_abs_pos_embed_nhwc, RotaryEmbeddingCat, apply_rot_embed_cat, to_2tuple, use_fused_attn
  from ._builder import build_model_with_cfg
  from ._manipulate import checkpoint_seq
  from ._registry import generate_default_cfgs, register_model
  from ._features_fx import register_notrace_function
label_desc: Use image node_id 3 for calling a global function with example usage: window_partition(x, window_size) and returns: windows,  Hp,  Wp