Convert x_pred network output to (v, x). Drops decoder_logits if present.

Return the shadow tensor for `name` on `p`'s device, migrating lazily. Shadow is built before Lightning moves the model; first call after the

Reshape attn_mask to (B, 1, Lq_or_1, Ls) for SDPA broadcasting. 2D (B, L) -> (B, 1, 1, L); 3D (B, L, L) -> (B, 1, L, L).

Build a zero tensor with cond positions restored from `cond_seq`.

Factory: resolve `name` to a registered encoder class and construct it.

Decode z (at t=1) to token ids via the decoder head.

Run a full sampling trajectory (SDE or ODE inner steps + final ODE step at t=1).

Mask everything at/after the first EOS per row (keeping EOS itself).

FlashAttention-4 path via `flash_attn.cute.flash_attn_func` (CuTeDSL, optimized for Hopper/Blackwell). Q/K/V cast to bf16 at the boundary; output

PyTorch SDPA in bf16 — kept as the masked-attention fallback because FA4's `flash_attn_func` does not accept an arbitrary dense (B, L, L) mask

Newton-Schulz iteration to compute the orthogonalization of G. Iterates X_{k+1} = a*X_k + (b*X_k X_k^T + c*(X_k X_k^T)^2) X_k for tuned (a,b,c).

Split params into (Muon, AdamW) groups: 2D weights → Muon; everything else (biases, RMSNorm gains, 1D learned tokens) → AdamW.

Return (encoder_attention_mask, attention_mask, cond_seq_mask) as float32 arrays. encoder_attention_mask: (B, L, L) — cond rows attend only to co

Run the frozen encoder and normalize the latents.

Return last_hidden_state of shape (B, L, d_model).