dew.nn.diffusion_gemma

DiffusionGemma’s shared native encoder/decoder and self-conditioning MLP.

The MLP follows Transformers modeling_diffusion_gemma.py:790-823: a scaled pre-norm, gated feed-forward, and scale-free post-norm. Previous logits become soft embeddings through an fp32 softmax against the scaled embedding table. An explicit self-conditioning mask zeros embeddings for the first inference step. The official SFT objective instead supplies zero logits for its dropout branch; those are uniform soft embeddings, not a zero signal.

Name	Summary
SelfConditioning	The previous step’s soft embeddings folded into the canvas embeddings.
soft_embeddings	Previous logits as soft embeddings: fp32 softmax against the table.
DiffusionGemma	One text parameter tree, read causally for context and bidirectionally for canvases.
translate_weights	Self-conditioning parameters, cast per weight before the layout copy.

SelfConditioning

Flax module source

class SelfConditioning(
    norm_eps: float = 1e-06,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
)

The previous step’s soft embeddings folded into the canvas embeddings.

SelfConditioning.setup

def setup()

soft_embeddings

function source

def soft_embeddings(
    logits: jax.typing.ArrayLike,
    embed_weight: jax.typing.ArrayLike,
    scale: float,
) -> jax.Array

Previous logits as soft embeddings: fp32 softmax against the table.

The table is contracted in its stored dtype with fp32 accumulation, which is the upcast product up to summation order and materialises no fp32 copy of the vocabulary-sized table.

DiffusionGemma

Flax module source

class DiffusionGemma(conditioner: VisionConditioner | None = None)

One text parameter tree, read causally for context and bidirectionally for canvases.

encode appends clean tokens to the cache. __call__ refines a canvas against that frozen cache and feeds previous logits through self-conditioning. Each method is a separate apply: sharing scopes keeps the encoder and decoder parameters identical without storing a second tree.

DiffusionGemma.setup

def setup()

DiffusionGemma.init_cache

def init_cache(batch_size: int)

DiffusionGemma.encode

def encode(
    tokens,
    *,
    positions=None,
    segment_ids=None,
    image_indices=None,
    attention_mask=None,
    image_groups=None,
    rotary_positions=None,
    attention_pairwise_mask=None,
    attention_key_positions=None,
    conditioning: Mapping[str, jax.Array] | None = None,
    train: bool = False,
    states: bool = False,
)

Append a clean prompt or committed canvas, evaluating media only when supplied.

The logits, or with states the final normalized states before the head: what a loss that scores the vocabulary a tile at a time reads, so the vocabulary-sized logits of a whole row never exist at once.

DiffusionGemma.head_weight

def head_weight(params)

The [D, vocab] head the encoder and the decoder score with, from the text tree of params, in its stored dtype (CausalTransformer.head_weight).

DiffusionGemma.head_table

def head_table(params)

The head as the text tree stores it and whether its rows are the vocabulary (CausalTransformer.head_table).

translate_weights

function source

def translate_weights(
    hf_tensors: Mapping[str, np.ndarray],
    *,
    param_dtype: str = 'float32',
) -> dict[str, dict[str, np.ndarray]]

Self-conditioning parameters, cast per weight before the layout copy.