Skip to navigation
DevelopmentAPI ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsStep5nemo_automodel.components.models.step5.vision

nemo_automodel.components.models.step5.vision

View as Markdown

Step perception encoder with shared, configuration-derived rotary tables.

Module Contents

Classes

NameDescription
Step5VisionEncoderReuse the Step encoder while owning derived-state initialization.
Step5VisionRotaryShare one derived table across vision layers; never serialize or cast it.

API

class nemo_automodel.components.models.step5.vision.Step5VisionEncoder(
    config: nemo_automodel.components.models.step5.configuration.Step5VisionConfig
)

Bases: StepRoboticsVisionEncoder

Reuse the Step encoder while owning derived-state initialization.

nemo_automodel.components.models.step5.vision.Step5VisionEncoder.forward(
    pixels: torch.Tensor,
    downsample: bool = False
) -> torch.Tensor

Encode images, optionally completing projection inputs inside the FSDP owner.

Parameters:

pixels
torch.Tensor

Images of shape [images, channels, height, width].

downsample
boolDefaults to False

Apply both checkpoint-owned stride-2 convolutions before FSDP reshards the vision module’s parameters.

Returns: torch.Tensor

Patch features [images, patches, width], or downsampled features

nemo_automodel.components.models.step5.vision.Step5VisionEncoder.reset_parameters() -> None

Initialize every vision parameter after allocation from meta.

class nemo_automodel.components.models.step5.vision.Step5VisionRotary(
    config: nemo_automodel.components.models.step5.configuration.Step5VisionConfig
)

Bases: Module

Share one derived table across vision layers; never serialize or cast it.

_frequencies
Tensor | None = None
_signature
tuple[device, dtype, int, int] | None = None
dim
= config.width // config.heads
maximum_grid
= config.image_size // config.patch_size
theta
use_cls_token
= config.use_cls_token
nemo_automodel.components.models.step5.vision.Step5VisionRotary.forward(
    q: torch.Tensor,
    k: torch.Tensor,
    grid_hw: tuple[int, int]
) -> tuple[torch.Tensor, torch.Tensor]

Apply the source encoder’s interleaved 2D RoPE.

Parameters:

q
torch.Tensor

Tensor [images, heads, image_tokens, head_dim].

k
torch.Tensor

Tensor with q’s shape and dtype.

grid_hw
tuple[int, int]

Patch-grid height and width; image_tokens includes the optional leading class token.

Returns: torch.Tensor

Rotated Q and K with their original shape and dtype. Frequencies