def get_2d_sine_cosine_positional_embedding(
    embed_dim: int, grid_size: int, cls_token: bool
) -> NDArray[np.float32]:
    """Generates 2D sine-cosine positional embedding.

    Code follows: https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py

    Args:
        embed_dim:
            Embedding dimension.
        grid_size:
            Height and width of the grid.
        cls_token:
            If True, a positional embedding for the class token is generated.
    Returns:
        Positional embedding with shape (grid_size * grid_size, embed_dim) or
        (1 + grid_size * grid_size, embed_dim) if cls_token is True.
    """
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = get_2d_sine_cosine_positional_embedding_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed
Generates 2D sine-cosine positional embedding.

    Code follows: https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py

    Args:
        embed_dim:
            Embedding dimension.
        grid_size:
            Height and width of the grid.
        cls_token:
            If True, a positional embedding for the class token is generated.
    Returns:
        Positional embedding with shape (grid_size * grid_size, embed_dim) or
        (1 + grid_size * grid_size, embed_dim) if cls_token is True.
    
lightly/models/utils.py