nemo_automodel.components.models.hunyuan_image3.rope

2D rotary position embedding of HunyuanImage-3.0.

Every token gets a (y, x) position. Text tokens use their sequence index for both coordinates, so a text-only sequence reduces to ordinary 1D RoPE. An image of h x w tokens starting at sequence index L is placed on a grid centered inside the span it occupies: y = L + (h*w - h) / 2 + row and x = L + (h*w - w) / 2 + col. Tokens after the image continue from L + h*w.

The rotary frequencies alternate between the two axes: of the head_dim / 2 frequencies, even ones rotate with y and odd ones with x. The angles are laid out for the half-split rotate_half convention.

Module Contents

Functions

Name	Description
`_rotate_half`	-
`apply_rope`	Rotate `x` of shape `[batch, heads, seq, head_dim]` with `[batch, seq, head_dim]` tables.
`image_grid_positions_batched`	Per-token (y, x) positions of a batch of sequences that each hold one image.
`rope_cos_sin`	Compute fp32 rotary tables from (y, x) positions.
`text_positions`	Return `[seq_len, 2]` positions of a text-only sequence (y = x = index).

API

nemo_automodel.components.models.hunyuan_image3.rope._rotate_half(
    x: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.hunyuan_image3.rope.apply_rope(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

Rotate x of shape [batch, heads, seq, head_dim] with [batch, seq, head_dim] tables.

The result is fp32 (the tables are fp32), matching the reference, which rotates and then normalizes q / k in fp32 before casting back.

nemo_automodel.components.models.hunyuan_image3.rope.image_grid_positions_batched(
    seq_len: int,
    image_starts: torch.Tensor,
    token_h: int,
    token_w: int
) -> torch.Tensor

Per-token (y, x) positions of a batch of sequences that each hold one image.

Parameters:

seq_len

int

Padded sequence length.

image_starts

torch.Tensor

Long tensor of shape [batch], index of the first image token of every sample; every image span must fit in seq_len.

token_h

int

Image height in tokens.

token_w

int

Image width in tokens.

Returns: torch.Tensor

Long tensor of shape [batch, seq_len, 2] with the (y, x) position of every token.

nemo_automodel.components.models.hunyuan_image3.rope.rope_cos_sin(
    positions: torch.Tensor,
    head_dim: int,
    base: float = 10000.0
) -> tuple[torch.Tensor, torch.Tensor]

Compute fp32 rotary tables from (y, x) positions.

Parameters:

positions

torch.Tensor

Integer tensor [..., seq, 2] of (y, x) positions.

head_dim

int

Attention head dimension; must be divisible by 4.

base

floatDefaults to 10000.0

RoPE base frequency.

Returns: tuple[torch.Tensor, torch.Tensor]

(cos, sin), each fp32 of shape [..., seq, head_dim].

nemo_automodel.components.models.hunyuan_image3.rope.text_positions(
    seq_len: int,
    device: torch.device | None = None
) -> torch.Tensor

Return [seq_len, 2] positions of a text-only sequence (y = x = index).