nemo_automodel.components.models.bagel.embeddings

Positional and timestep embeddings for BAGEL modules.

Module Contents

Classes

Name	Description
`BagelGridPositionEmbedding`	Frozen 2D sine/cosine position table for patch and latent grids.
`BagelTimestepEmbedding`	Map scalar timesteps through sine/cosine features and a small MLP.

Functions

Name	Description
`_build_2d_sincos_pos_embed`	Return `(grid_size ** 2, embed_dim)` sine/cosine features for a square grid.
`_build_square_grid_positions`	-
`_encode_scalar_positions`	Encode scalar coordinates with sine features followed by cosine features.
`_geometric_frequencies`	Return inverse periods for one sine/cosine axis.
`_get_2d_sincos_pos_embed`	Return a 2D sin-cos embedding table as a NumPy array.
`_timestep_features`	Create sine/cosine features for scalar timesteps.
`_validate_even`	-
`_validate_positive`	-

Data

_SIN_COS_BASE

API

class nemo_automodel.components.models.bagel.embeddings.BagelGridPositionEmbedding(
    max_num_patch_per_side: int,
    hidden_size: int
)

Bases: Module

Frozen 2D sine/cosine position table for patch and latent grids.

pos_embed

nemo_automodel.components.models.bagel.embeddings.BagelGridPositionEmbedding._init_weights() -> None

nemo_automodel.components.models.bagel.embeddings.BagelGridPositionEmbedding.forward(
    position_ids: torch.Tensor
) -> torch.Tensor

class nemo_automodel.components.models.bagel.embeddings.BagelTimestepEmbedding(
    hidden_size: int,
    frequency_embedding_size: int = 256
)

Bases: Module

Map scalar timesteps through sine/cosine features and a small MLP.

mlp

nemo_automodel.components.models.bagel.embeddings.BagelTimestepEmbedding.forward(
    t: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.bagel.embeddings._build_2d_sincos_pos_embed(
    embed_dim: int,
    grid_size: int
) -> torch.Tensor

Return (grid_size ** 2, embed_dim) sine/cosine features for a square grid.

nemo_automodel.components.models.bagel.embeddings._build_square_grid_positions(
    grid_size: int
) -> tuple[torch.Tensor, torch.Tensor]

nemo_automodel.components.models.bagel.embeddings._encode_scalar_positions(
    positions: torch.Tensor,
    width: int
) -> torch.Tensor

Encode scalar coordinates with sine features followed by cosine features.

nemo_automodel.components.models.bagel.embeddings._geometric_frequencies(
    width: int,
    device: torch.device | None = None
) -> torch.Tensor

Return inverse periods for one sine/cosine axis.

nemo_automodel.components.models.bagel.embeddings._get_2d_sincos_pos_embed(
    embed_dim: int,
    grid_size: int
) -> numpy.ndarray

Return a 2D sin-cos embedding table as a NumPy array.

nemo_automodel.components.models.bagel.embeddings._timestep_features(
    timesteps: torch.Tensor,
    width: int,
    max_period: int = 10000
) -> torch.Tensor

Create sine/cosine features for scalar timesteps.

nemo_automodel.components.models.bagel.embeddings._validate_even(
    value: int,
    name: str
) -> None

nemo_automodel.components.models.bagel.embeddings._validate_positive(
    value: int,
    name: str
) -> None

nemo_automodel.components.models.bagel.embeddings._SIN_COS_BASE = 10000.0