nemo_automodel.components.models.inkling2.vision

Native Inkling2 image encoder with independent per-image SDPA attention.

The tensor contracts follow the released Inkling-2-Small checkpoint configuration. This implementation uses ordinary Torch modules and SDPA, preserving the released weight names for bounded checkpoint loading.

Module Contents

Classes

Name	Description
`Inkling2VisionAttention`	Full attention with separate horizontal and vertical rotary coordinates.
`Inkling2VisionBackbone`	Image transformer before spatial merging and text projection.
`Inkling2VisionEmbeddings`	Native patch projection, interpolated positions and a learned class token.
`Inkling2VisionEncoder`	Encode merged image patches into one text feature per image-body token.
`Inkling2VisionLayer`	Pre-LayerNorm attention and GELU branches with learned LayerScale.
`Inkling2VisionLayerScale`	Learned channel scale on a vision residual branch.
`Inkling2VisionMLP`	The image encoder’s ordinary GELU feed-forward block.

API

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionAttention(
    config: nemo_automodel.components.models.inkling2.config.Inkling2VisionConfig,
    dtype: torch.dtype
)

Bases: Module

Full attention with separate horizontal and vertical rotary coordinates.

head_dim

= config.hidden_size // self.num_heads

k_proj

num_heads

= config.num_attention_heads

out_proj

q_proj

v_proj

nemo_automodel.components.models.inkling2.vision.Inkling2VisionAttention.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor
) -> torch.Tensor

Attend over one image, including CLS; never mix different images.

Parameters:

hidden

torch.Tensor

Compute-dtype [tokens, hidden] features.

positions

torch.Tensor

Integer [tokens, 2] (y, x) coordinates; CLS uses zero.

Returns: torch.Tensor

Projected features with the same shape and dtype as hidden.

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionBackbone(
    config: nemo_automodel.components.models.inkling2.config.Inkling2VisionConfig,
    dtype: torch.dtype
)

Bases: Module

Image transformer before spatial merging and text projection.

embeddings

= Inkling2VisionEmbeddings(config, dtype=dtype)

layers

ln_pre

nemo_automodel.components.models.inkling2.vision.Inkling2VisionBackbone.forward(
    pixels: torch.Tensor,
    height: int,
    width: int
) -> torch.Tensor

Encode one image’s native patches.

Parameters:

pixels

torch.Tensor

[height * width, native_patch_size ** 2 * 3] rows in native row-major patch order, as consumed by the embeddings.

height

int

Native patch rows of the image.

width

int

Native patch columns of the image.

Returns: torch.Tensor

New [height * width, hidden] features in the same row-major patch order, without

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionEmbeddings(
    config: nemo_automodel.components.models.inkling2.config.Inkling2VisionConfig,
    dtype: torch.dtype
)

Bases: Module

Native patch projection, interpolated positions and a learned class token.

class_embedding

= nn.Parameter(torch.empty(width, dtype=dtype))

cls_position_embedding

= nn.Parameter(torch.empty(width, dtype=dtype))

grid_size

= config.posemb_grid_size

patch_embedding

position_embedding

nemo_automodel.components.models.inkling2.vision.Inkling2VisionEmbeddings.forward(
    pixels: torch.Tensor,
    height: int,
    width: int
) -> torch.Tensor

Embed one image’s native patches and add interpolated positions.

Parameters:

pixels

torch.Tensor

Tensor of shape [height * width, native_patch_size ** 2 * 3]. Rows follow the native patch grid in row-major order; each row flattens patch height, patch width and RGB channels, in that order. Its device and dtype must match the patch projection.

height

int

Native patch rows of the image.

width

int

Native patch columns of the image.

Returns: torch.Tensor

New [1 + height * width, hidden] tokens in the projection dtype, with the class

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionEncoder(
    config: nemo_automodel.components.models.inkling2.config.Inkling2VisionConfig,
    text_hidden_size: int,
    dtype: torch.dtype
)

Bases: Module

Encode merged image patches into one text feature per image-body token.

_clip_mean

Tensor

_clip_std

Tensor

backbone

= Inkling2VisionBackbone(config, dtype=dtype)

image_cp_mesh

DeviceMesh | None = None

proj_in

proj_out

nemo_automodel.components.models.inkling2.vision.Inkling2VisionEncoder.forward(
    vision_patches_bthwc: torch.Tensor,
    vision_grid_thw: torch.Tensor
) -> torch.Tensor

Project independent images into text embeddings.

Parameters:

vision_patches_bthwc

torch.Tensor

BF16 normalized patches [patches, temporal, side, side, 3], where side = native_patch_size * adapter_merge_size. Input is unchanged.

vision_grid_thw

torch.Tensor

Integer [images, 3] grid metadata (1, rows, columns). Grids and patches follow the same image order; metadata may be on CPU or GPU.

Returns: torch.Tensor

Compute-dtype [patches, text_hidden] features, in row-major merged-patch order.

nemo_automodel.components.models.inkling2.vision.Inkling2VisionEncoder.reset_parameters() -> None

Initialize all parameters and regenerate nonpersistent normalization buffers.

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionLayer(
    config: nemo_automodel.components.models.inkling2.config.Inkling2VisionConfig,
    dtype: torch.dtype
)

Bases: Module

Pre-LayerNorm attention and GELU branches with learned LayerScale.

ln_1

ln_2

ls_1

ls_2

mlp

= Inkling2VisionMLP(config, dtype=dtype)

self_attn

= Inkling2VisionAttention(config, dtype=dtype)

nemo_automodel.components.models.inkling2.vision.Inkling2VisionLayer.forward(
    hidden: torch.Tensor,
    positions: torch.Tensor
) -> torch.Tensor

Apply one pre-LayerNorm attention and feed-forward block.

Parameters:

hidden

torch.Tensor

[tokens, hidden] features of one image, class token first.

positions

torch.Tensor

Integer [tokens, 2] (y, x) patch coordinates; the class token uses zero.

Returns: torch.Tensor

New [tokens, hidden] features; the inputs are unchanged.

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionLayerScale(
    width: int,
    dtype: torch.dtype
)

Bases: Module

Learned channel scale on a vision residual branch.

gamma

= nn.Parameter(torch.ones(width, dtype=dtype))

nemo_automodel.components.models.inkling2.vision.Inkling2VisionLayerScale.forward(
    hidden: torch.Tensor
) -> torch.Tensor

Scale a residual branch per channel.

Parameters:

hidden

torch.Tensor

[tokens, hidden] branch output.

Returns: torch.Tensor

New [tokens, hidden] features, each channel multiplied by the learned [hidden]

class nemo_automodel.components.models.inkling2.vision.Inkling2VisionMLP(
    config: nemo_automodel.components.models.inkling2.config.Inkling2VisionConfig,
    dtype: torch.dtype
)

Bases: Module

The image encoder’s ordinary GELU feed-forward block.

c_fc

c_proj

nemo_automodel.components.models.inkling2.vision.Inkling2VisionMLP.forward(
    hidden: torch.Tensor
) -> torch.Tensor

Apply the GELU feed-forward block.

Parameters:

hidden

torch.Tensor

[tokens, hidden] features in the block’s dtype.

Returns: torch.Tensor

New [tokens, hidden] features; the input is unchanged.