`nemo_automodel.components.models.bagel.embeddings`#

Positional and timestep embeddings for BAGEL modules.

Module Contents#

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

`_validate_positive`
`_validate_even`
`_geometric_frequencies`	Return inverse periods for one sine/cosine axis.
`_encode_scalar_positions`	Encode scalar coordinates with sine features followed by cosine features.
`_build_square_grid_positions`
`_build_2d_sincos_pos_embed`	Return `(grid_size ** 2, embed_dim)` sine/cosine features for a square grid.
`_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.

nemo_automodel.components.models.bagel.embeddings._validate_positive(value: int, *, name: str) → None#

nemo_automodel.components.models.bagel.embeddings._validate_even(value: int, *, name: str) → None#

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._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._build_square_grid_positions( grid_size: int, ) → tuple[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._get_2d_sincos_pos_embed( embed_dim: int, grid_size: int, ) → numpy.ndarray#: Return a 2D sin-cos embedding table as a NumPy array.

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

Bases: torch.nn.Module

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

Initialization

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.

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

Bases: torch.nn.Module

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

Initialization