Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsPraxisnemo_automodel.components.models.praxis.backbone

nemo_automodel.components.models.praxis.backbone

View as MarkdownOpen in Claude

Two-tower video backbone of the chunk-causal window.

The encoder tower reads [registers | anchor | text | noised chunks | clean chunks] and returns the conditioning of the noised tokens plus the tap features for the action expert. The decoder tower reads [registers | noised chunks] at the unpatched latent resolution, modulated per token by the encoder output, and returns the video velocity.

Rotary positions: each latent frame owns a block of FRAME_STRIDE positions; inside a frame a token sits at 1 + row * ROW_STRIDE + column. Registers are at position 0. In the encoder the anchor takes frame block 0, the text frame block 1 and the noised frames start at frame block 2; a clean frame shares the positions of its noised twin.

Module Contents

Classes

NameDescription
BackboneOutputvelocity of shape [batch, channels, frames, height, width] (None without decoding) and the encoder features.
DecoderOutputProjectionZero-centred RMSNorm followed by a linear map to latent channels.
DecoderTowerThe decoder tower (see the module docstring).
EncoderFeaturesEncoder outputs that the action expert reads.
EncoderTowerThe encoder tower (see the module docstring).
FourierTimestepFixed random Fourier features of the flow time: 0.5 * cos(2 pi (t f + phi)).
MoeTrainingSettingsLoad-balancing terms of the routers during training (both zero at inference).
TextInputProjectionMLP from text-encoder states to trunk tokens, ending in a zero-centred RMSNorm.
VideoBackboneEncoder tower and decoder tower of one denoising step.

Functions

NameDescription
_expand_rowsRepeat row k of rows ([batch, chunks, hidden]) over chunk_tokens[k] tokens.
_make_stack-
encoder_positionsRotary positions of the encoder sequence [registers | anchor | text | noised | clean].
patchifyInverse of unpatch: fold (ph, pw) spatial patches into the channels.
unpatchMove the channel patch of [batch, channels * ph * pw, frames, h, w] into space.
video_positionsRotary positions of a video run in frame-major order.

Data

FRAME_STRIDE

ROW_STRIDE

API

class nemo_automodel.components.models.praxis.backbone.BackboneOutput(
    velocity: torch.Tensor | None,
    features: nemo_automodel.components.models.praxis.backbone.EncoderFeatures
)
Dataclass

velocity of shape [batch, channels, frames, height, width] (None without decoding) and the encoder features.

features
EncoderFeatures
velocity
Tensor | None
class nemo_automodel.components.models.praxis.backbone.DecoderOutputProjection(
    dim: int,
    c_out: int
)

Bases: Module

Zero-centred RMSNorm followed by a linear map to latent channels.

linear
= nn.Linear(dim, c_out)
norm
= ZeroCenteredRMSNorm(dim, cast_dtype=None)
nemo_automodel.components.models.praxis.backbone.DecoderOutputProjection.forward(
    x: torch.Tensor
) -> torch.Tensor

Project decoder tokens.

Parameters:

x
Tensor

Tensor of shape [batch, tokens, hidden].

Returns: Tensor

Tensor of shape [batch, tokens, c_out].

Bases: Module

The decoder tower (see the module docstring).

_masks
= MaskCache()
block_size
= cfg.attention.flex_block_size
core
modulation_dtype
= cfg.modulation_dtype
patch_size
= (tower.unpatch_factor[1], tower.unpatch_factor[2])
registers
t_embedding
= FourierTimestep(cfg.timestep_dim)
time_proj
= nn.Linear(cfg.timestep_dim, 2 * cfg.dim)
use_flex
= backend.attn == 'flex'
video_in
= nn.Linear(tower.video_in.c_in, cfg.dim)
video_out
= DecoderOutputProjection(cfg.dim, c_out)
nemo_automodel.components.models.praxis.backbone.DecoderTower._mask(
    slots: tuple[nemo_automodel.components.models.praxis.masks.Slot, ...],
    chunk: int,
    device: torch.device
) -> torch.nn.attention.flex_attention.BlockMask | torch.Tensor
nemo_automodel.components.models.praxis.backbone.DecoderTower.forward(
    noised: torch.Tensor,
    cond: torch.Tensor,
    timesteps: torch.Tensor,
    rope_shift: int = 0
) -> torch.Tensor

Predict the velocity of the noised chunks.

Parameters:

noised
Tensor

Noised latents of shape [batch, channels, K * F, h, w].

cond
Tensor

Encoder conditioning of shape [batch, K * F, h, w, hidden].

timesteps
Tensor

Flow time of each chunk, shape [batch, K].

rope_shift
intDefaults to 0

Rotary offset of the video positions in training.

Returns: Tensor

Velocity of shape [batch, channels, K * F, h, w].

class nemo_automodel.components.models.praxis.backbone.EncoderFeatures(
    taps: list[torch.Tensor],
    embeddings: torch.Tensor,
    positions: torch.Tensor,
    global_spans: dict[str, tuple[int, int]],
    noised_span: tuple[int, int],
    clean_span: tuple[int, int] | None,
    frames_per_chunk: int,
    tokens_per_frame: int,
    n_chunks: int
)
Dataclass

Encoder outputs that the action expert reads.

clean_span
tuple[int, int] | None

(start, length) of the clean chunk tokens, None for one chunk.

embeddings
Tensor

The encoder input sequence (block 0 input) of shape [batch, tokens, hidden].

frames_per_chunk
int

Latent frames per chunk.

global_spans
dict[str, tuple[int, int]]

(start, length) of the register, anchor and text runs.

n_chunks
int

Chunks in the window.

noised_span
tuple[int, int]

(start, length) of the noised chunk tokens.

positions
Tensor

Rotary positions of the sequence, shape [batch, tokens].

taps
list[Tensor]

Block outputs at the tap blocks, each of shape [batch, tokens, hidden].

tokens_per_frame
int

Tokens per latent frame.

Bases: Module

The encoder tower (see the module docstring).

_masks
= MaskCache()
block_size
= cfg.attention.flex_block_size
core
modulation_dtype
= cfg.modulation_dtype
registers
t_embedding
= FourierTimestep(cfg.timestep_dim)
tap_blocks
= tower.tap_blocks
text_in
= TextInputProjection(tower.text_in.c_in, cfg.dim)
time_proj
= nn.Linear(cfg.timestep_dim, 2 * cfg.dim)
use_flex
= backend.attn == 'flex'
video_in
= nn.Linear(tower.video_in.c_in, cfg.dim)
video_out
= nn.Linear(cfg.dim, cfg.dim)
nemo_automodel.components.models.praxis.backbone.EncoderTower._mask(
    slots: tuple[nemo_automodel.components.models.praxis.masks.Slot, ...],
    chunk: int,
    device: torch.device
) -> torch.nn.attention.flex_attention.BlockMask | torch.Tensor
nemo_automodel.components.models.praxis.backbone.EncoderTower._modulate_chunks(
    tokens: torch.Tensor,
    timesteps: torch.Tensor,
    chunk_tokens: list[int]
) -> torch.Tensor
nemo_automodel.components.models.praxis.backbone.EncoderTower._video_tokens(
    latents: torch.Tensor
) -> torch.Tensor

[batch, channels, frames, h, w] latents to [batch, frames * h * w, hidden] tokens.

nemo_automodel.components.models.praxis.backbone.EncoderTower.forward(
    text: torch.Tensor,
    anchor: torch.Tensor,
    clean: torch.Tensor | None,
    noised: torch.Tensor,
    timesteps: torch.Tensor,
    rope_shifts: tuple[int, int, int] = (0, 0, 0)
) -> tuple[torch.Tensor, nemo_automodel.components.models.praxis.backbone.EncoderFeatures]

Encode a window of K chunks.

Parameters:

text
Tensor

Text-encoder states of shape [batch, text_tokens, text_dim].

anchor
Tensor

Anchor latent of shape [batch, channels, 1, h, w].

clean
Tensor | None

Clean copies of chunks 0 .. K-2 of shape [batch, channels, (K-1) * F, h, w], or None for K == 1.

noised
Tensor

Noised chunks of shape [batch, channels, K * F, h, w].

timesteps
Tensor

Flow time of each noised chunk, shape [batch, K].

rope_shifts
tuple[int, int, int]Defaults to (0, 0, 0)

(anchor, text, video) rotary offsets of training.

Returns: Tensor

The conditioning of the noised tokens, shape [batch, K * F, h, w, hidden], and the

class nemo_automodel.components.models.praxis.backbone.FourierTimestep(
    dim: int
)

Bases: Module

Fixed random Fourier features of the flow time: 0.5 * cos(2 pi (t f + phi)).

freqs
Tensor
phases
Tensor
nemo_automodel.components.models.praxis.backbone.FourierTimestep.forward(
    t: torch.Tensor
) -> torch.Tensor

Embed timesteps.

Parameters:

t
Tensor

Tensor of shape […, 1] or […]; the last axis is broadcast against the features.

Returns: Tensor

Tensor of shape […, dim] in the dtype of t.

class nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings(
    gate_bias_update_factor: float = 0.0,
    aux_loss_coeff: float = 0.0
)
Dataclass

Load-balancing terms of the routers during training (both zero at inference).

aux_loss_coeff
float = 0.0

Weight of the sequence-wise balance loss of each router (experts * sum_i f_i p_i), injected through the routing weights.

gate_bias_update_factor
float = 0.0

Step of the per-expert selection bias toward balanced load, applied by PraxisPolicy.update_moe_gate_bias after each optimizer step.

class nemo_automodel.components.models.praxis.backbone.TextInputProjection(
    c_in: int,
    dim: int
)

Bases: Module

MLP from text-encoder states to trunk tokens, ending in a zero-centred RMSNorm.

fc1
= nn.Linear(c_in, 4 * dim)
fc2
= nn.Linear(4 * dim, dim)
norm
= ZeroCenteredRMSNorm(dim, cast_dtype=None)
nemo_automodel.components.models.praxis.backbone.TextInputProjection.forward(
    x: torch.Tensor
) -> torch.Tensor

Project text states.

Parameters:

x
Tensor

Tensor of shape [batch, text_tokens, c_in].

Returns: Tensor

Tensor of shape [batch, text_tokens, hidden].

class nemo_automodel.components.models.praxis.backbone.VideoBackbone(
    cfg: nemo_automodel.components.models.praxis.config.BackboneConfig,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    dtype: torch.dtype,
    moe_training: nemo_automodel.components.models.praxis.backbone.MoeTrainingSettings | None = None
)

Bases: Module

Encoder tower and decoder tower of one denoising step.

decoder
= DecoderTower(cfg, backend, dtype, moe_training)
encoder
= EncoderTower(cfg, backend, dtype, moe_training)
nemo_automodel.components.models.praxis.backbone.VideoBackbone.forward(
    text: torch.Tensor,
    anchor: torch.Tensor,
    clean: torch.Tensor | None,
    noised: torch.Tensor,
    timesteps: torch.Tensor,
    decode: bool = True,
    rope_shifts: tuple[int, int, int, int] = (0, 0, 0, 0)
) -> nemo_automodel.components.models.praxis.backbone.BackboneOutput

Run one denoising step of a window.

Parameters:

text
Tensor

Text-encoder states of shape [batch, text_tokens, text_dim].

anchor
Tensor

Anchor latent of shape [batch, channels, 1, h, w].

clean
Tensor | None

Clean copies of shape [batch, channels, (K-1) * F, h, w], or None.

noised
Tensor

Noised chunks of shape [batch, channels, K * F, h, w].

timesteps
Tensor

Flow time of each chunk, shape [batch, K].

decode
boolDefaults to True

Run the decoder tower; otherwise velocity is None.

rope_shifts
tuple[int, int, int, int]Defaults to (0, 0, 0, 0)

(encoder anchor, encoder text, encoder video, decoder video) offsets.

Returns: BackboneOutput

class:BackboneOutput with the velocity of shape [batch, channels, K * F, h, w].

nemo_automodel.components.models.praxis.backbone._expand_rows(
    rows: torch.Tensor,
    chunk_tokens: list[int]
) -> torch.Tensor

Repeat row k of rows ([batch, chunks, hidden]) over chunk_tokens[k] tokens.

 
nemo_automodel.components.models.praxis.backbone.encoder_positions(
    num_registers: int,
    anchor_hw: tuple[int, int],
    num_text: int,
    noised_thw: tuple[int, int, int],
    num_clean_frames: int,
    shifts: tuple[int, int, int]
) -> torch.Tensor

Rotary positions of the encoder sequence [registers | anchor | text | noised | clean].

Parameters:

num_registers
int

Register tokens (position 0).

anchor_hw
tuple[int, int]

Latent grid of the anchor frame.

num_text
int

Text tokens.

noised_thw
tuple[int, int, int]

Frames and latent grid of the noised chunks.

num_clean_frames
int

Frames of the clean copies.

shifts
tuple[int, int, int]

Random (anchor, text, video) offsets of training; zeros at inference.

Returns: Tensor

Int64 tensor of shape [tokens].

nemo_automodel.components.models.praxis.backbone.patchify(
    x: torch.Tensor,
    patch: tuple[int, int]
) -> torch.Tensor

Inverse of unpatch: fold (ph, pw) spatial patches into the channels.

Returns: Tensor

Tensor of shape [batch, channels * ph * pw, frames, h // ph, w // pw].

nemo_automodel.components.models.praxis.backbone.unpatch(
    x: torch.Tensor,
    patch: tuple[int, int]
) -> torch.Tensor

Move the channel patch of [batch, channels * ph * pw, frames, h, w] into space.

Returns: Tensor

Tensor of shape [batch, channels, frames, h * ph, w * pw].

nemo_automodel.components.models.praxis.backbone.video_positions(
    frames: int,
    height: int,
    width: int,
    start: int
) -> torch.Tensor

Rotary positions of a video run in frame-major order.

Returns: Tensor

Int64 tensor of shape [frames * height * width].

nemo_automodel.components.models.praxis.backbone.FRAME_STRIDE = 129600
 
nemo_automodel.components.models.praxis.backbone.ROW_STRIDE = 360