def get_2d_sine_cosine_positional_embedding_from_grid(
    embed_dim: int, grid: NDArray[np.float32]
) -> NDArray[np.float32]:
    """Generates 2D sine-cosine positional embedding from a grid.

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

    Args:
        embed_dim:
            Embedding dimension.
        grid:
            Grid of shape (2, grid_size, grid_size) with x and y coordinates.
    Returns:
        Positional embedding with shape (grid_size * grid_size, embed_dim).
    """
    assert embed_dim % 2 == 0

    # Use half of dimensions to encode grid_h
    # (grid_size * grid_size, embed_dim/2)
    emb_h = get_1d_sine_cosine_positional_embedding_from_positions(
        embed_dim // 2, grid[0]
    )

    # Use the other half of dimensions to encode grid_w
    # (grid_size * grid_size, embed_dim/2)
    emb_w = get_1d_sine_cosine_positional_embedding_from_positions(
        embed_dim // 2, grid[1]
    )

    emb = np.concatenate([emb_h, emb_w], axis=1)  # (grid_size * grid_size, embed_dim)
    return emb
Generates 2D sine-cosine positional embedding from a grid.

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

    Args:
        embed_dim:
            Embedding dimension.
        grid:
            Grid of shape (2, grid_size, grid_size) with x and y coordinates.
    Returns:
        Positional embedding with shape (grid_size * grid_size, embed_dim).
    
lightly/models/utils.py