dew.nn.multimodal

Native vision conditioning around the shared causal decoder.

Name	Summary
Fusion	Decoder token identities, their scaled, conditioned embeddings, and which positions media fill.
VisionConditioner	An image encoder and projector with row-aligned media inputs.
AudioConditioner	An audio encoder and projector over row-aligned clips.
MultimodalTransformer	A media-conditioned decoder using the ordinary decoder cache and head.

Fusion

dataclass source

class Fusion()

Decoder token identities, their scaled, conditioned embeddings, and which positions media fill.

VisionConditioner

Flax module source

class VisionConditioner(dtype: Dtype | None = None, precision: PrecisionLike = None)

An image encoder and projector with row-aligned media inputs.

Pixels have shape [batch, images, channels, height, width]. All images in a numeric batch share their processed resolution; the processor keeps per-image lengths and does not substitute preprocessing inside the model.

VisionConditioner.setup

def setup()

VisionConditioner.initialize_parameters

def initialize_parameters() -> None

Create the media leaves during an ordinary token-only model init.

Fixed-resolution towers use their configured image size. Other towers need only one pooling/merge block to create resolution-independent parameters; real batches supply their processed geometry later.

VisionConditioner.fuse

def fuse(
    tokens: jax.Array,
    embeddings: jax.Array,
    image_indices: jax.Array,
    conditioning: Mapping[str, jax.Array],
    train: bool = False,
) -> Fusion

Replace marked text slots with the corresponding soft feature.

AudioConditioner

Flax module source

class AudioConditioner(
    soft_tokens: int | None = None,
    padding_id: int | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

An audio encoder and projector over row-aligned clips.

Features have shape [batch, clips, frames, mel] with a True-for-valid frame mask. The table returned holds every clip’s encoded frames back to back, clips * capacity per row, where capacity is the encoded frame count, or soft_tokens when the source fixes the slot count per clip (Gemma 3n): there, padded frames and the slots past the encoded frames carry the embedder’s padding token, as modeling_gemma3n.py does.

AudioConditioner.setup

def setup()

AudioConditioner.initialize_parameters

def initialize_parameters() -> None

Create the audio leaves during a token-only model init.

MultimodalTransformer

Flax module source

class MultimodalTransformer(
    pad_token_id: int = 0,
    extra_placeholder_ids: tuple[int, ...] = (),
    audio: TowerBase | None = None,
    audio_projection: ProjectorBase | None = None,
    audio_soft_tokens: int | None = None,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

A media-conditioned decoder using the ordinary decoder cache and head.

image_indices and audio_indices identify the soft feature for each text slot, or -1 for text. Media are evaluated during conditioned forward/prefill calls; subsequent decoding reads the cached language states without rerunning the encoders. Parameters retain the existing language_model, tower and projector names; audio adds audio_tower and audio_projector. Gemma 3n also embeds its hard vision and audio vocabulary ranges through the embedders and keeps placeholder ids for its per-layer inputs, as modeling_gemma3n.py does.

MultimodalTransformer.setup

def setup()

MultimodalTransformer.mtp_hidden_states

def mtp_hidden_states(
    hidden,
    tokens,
    train: bool = False,
    positions=None,
    segment_ids=None,
    image_indices=None,
    conditioning=None,
    attention_mask=None,
    image_groups=None,
    rotary_positions=None,
)

Prediction layers over the same media embeddings as the main decoder.

MultimodalTransformer.mtp_logits

def mtp_logits(hidden, tokens, **kwargs={})

The shared language head over each media-aware prediction depth.

MultimodalTransformer.mtp_step

def mtp_step(
    hidden,
    tokens,
    *,
    image_indices=None,
    conditioning=None,
    input_embeddings=None,
    **kwargs={},
)

One candidate prediction step using the decoder’s independent MTP cache, returning its logits and its hidden state.

MultimodalTransformer.token_embeddings

def token_embeddings(tokens)

The decoder’s own table; a drawn token is text, never media.

MultimodalTransformer.init_mtp_cache

def init_mtp_cache(batch_size: int)

MultimodalTransformer.hidden_states

def hidden_states(
    tokens,
    train: bool = False,
    decode: bool = False,
    positions=None,
    segment_ids=None,
    image_indices=None,
    conditioning: Mapping[str, jax.Array] | None = None,
    attention_mask=None,
    image_groups=None,
    rotary_positions=None,
    audio_indices=None,
    routed_experts=None,
    routed=None,
)

Run the decoder over text with the media embeddings spliced in.

conditioning holds the towers’ payloads and image_indices and audio_indices say which token positions each one replaces. A decode step tracks the next position in the cache collection, because a media span advances a row by more than one token. routed_experts and routed replay a routing record and pass to the language model unchanged: engines record a row for every placeholder position too, so the [B, S, layers, top_k] layout is the text’s.

MultimodalTransformer.states_and_logits

def states_and_logits(tokens, **kwargs={})

The final hidden states and their logits from one media-aware forward.

MultimodalTransformer.states_and_logits_at

def states_and_logits_at(tokens, slots, **kwargs={})

The final hidden states, and the logits of one slot per row.

A prefill scores the position the first draw reads and no other; the head over every prompt position is what a long request allocates most of its transient memory for.

MultimodalTransformer.head_weight

def head_weight(params)

The decoder’s shared fp32 head matrix for chunked objective scoring.

MultimodalTransformer.init_cache

def init_cache(batch_size: int)

Allocate the nested language cache without evaluating media.