def window_partition(x, window_size):
    """
    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.transpose([0, 3, 1, 2]),
              paddle.to_tensor(
                  [0, int(pad_w), 0, int(pad_h)],
                  dtype='int32')).transpose([0, 2, 3, 1])
    Hp, Wp = H + pad_h, W + pad_w

    num_h, num_w = Hp // window_size, Wp // window_size

    x = x.reshape([B, num_h, window_size, num_w, window_size, C])
    windows = x.transpose([0, 1, 3, 2, 4, 5]).reshape(
        [-1, window_size, window_size, C])
    return windows, (Hp, Wp), (num_h, num_w)

    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
    
ppdet/modeling/backbones/transformer_utils.py