Index
layers
¶
RotaryPosEncoding(d_model: int, base: int = 10000, seq_dim: int = 1, partial_rotary_factor: float = 1.0)
¶
Unified RoPE supporting standard, Qwen, and NeoX variants.
Differences between variants are controlled by constructor args: - base: 10000 (standard/NeoX) or 1e6 (Qwen) - partial_rotary_factor: 1.0 (standard/Qwen) or <1.0 (NeoX)
MRotaryPosEncoding(d_model: int, base: int = 10000, seq_dim: int = 1, partial_rotary_factor: float = 1.0, mrope_section: Optional[Sequence[int]] = None)
¶
Bases: RotaryPosEncoding
RoPE with optional 3-channel (T, H, W) position ids.
Pass positions as (B, T) for text-only; as (3, B, T) for
multimodal — the channels are interleaved per mrope_section.
RMSNorm
¶
Bases: Module
Root-mean-square layer norm.
centered=True switches to the Qwen 3.5 / OLMo-2 convention:
multiplier is 1 + weight with weight init at 0 (numerically still
centered at 1, but lets HF checkpoints round-trip).
RMSNormGated
¶
Bases: Module
RMSNorm with a multiplicative silu(gate) after the weight.
Used inside the gated-delta-net token mixer (Qwen 3.5 linear-attention
layers). Weight is initialized to ones — matches HF Qwen3_5RMSNormGated.
FLTMLP
¶
Bases: MLP
ReLU^2 MLP with a fixed 4x expansion and no bias.
rms(x: jax.Array, eps: float) -> jax.Array
¶
RMS normalization without a learned scale, in float32. FLT normalizes every sublayer input and the loop boundary this way; RMSNorm has a scale.