gemma4_jax — Flax NNX port of google/gemma-4-12B (full multimodal)

A from-scratch, faithful Flax NNX implementation of Google's Gemma 4 12B "Unified" model — the text decoder and the encoder-free vision + audio embedders — plus a HuggingFace→NNX weight converter. Built because the model is not (yet) in the official google-deepmind/gemma JAX library.

• 🤗 HuggingFace: https://huggingface.co/mlnomad/gemma4-jax
• 🧬 Base model: google/gemma-4-12B

This repo hosts the port (code). It does not redistribute weights — the converter reads the official google/gemma-4-12B safetensors directly (accept the Gemma terms first). The weights remain under the Gemma Terms of Use.

Status

Verified — param count is exact to the parameter. Text-only analytic = 11,907,350,320; vision+audio embedders = 52,379,904; full-model total = 11,959,730,224, which equals the published checkpoint exactly (0 diff). The multimodal smoke test confirms the vision/audio projectors run, that soft tokens actually splice into placeholder positions, and that the converter loads the real vision_embedder / embed_vision / embed_audio tensors and produces finite, well-scaled features. The text smoke test confirms forward pass, causality (drift = 0), and softcap bounds.

Not yet verified: numerical equivalence against the reference HF forward pass. That requires downloading the gated weights (see below) and diffing logits. The architecture matches the transformers gemma4_unified source, but treat exact-numerics parity as unconfirmed until you run that diff.

Provenance

Architecture derived from the published config.json of google/gemma-4-12B (text_config) and the HuggingFace transformers gemma4_unified modeling source (transformers 5.10.0.dev0). Non-obvious details that differ from Gemma 2/3 and are baked in here:

• RMSNorm is plain x_normed * weight (NOT (1+w)); eps is inside the inverse-sqrt; fp32 internals.
• Attention scaling = 1.0 (not 1/sqrt(d)) — query magnitude is set by a per-head q_norm RMSNorm instead.
• k_eq_v drops v_proj on full/global layers only; V = the K projection output (pre-k_norm, pre-RoPE), then a scale-free v_norm, and V is not rotated.
• Proportional RoPE on full layers = a length-head_dim/2 inv_freq where only the first int(0.25·512//2)=64 entries are real frequencies (base 1e6, denom = 512) and the rest are zero (a NoPE tail).
• layer_scalar (=1.0 in the checkpoint) multiplies each layer's output.
• Input embeddings scaled by sqrt(hidden_size); output 30·tanh(logits/30); no attention-logit softcapping.

Usage

Self-contained (recommended) — JAX-native weights, no PyTorch

import jax.numpy as jnp
from gemma4_jax.convert import unified_from_pretrained

# downloads JAX-native sharded weights from HF — no PyTorch, no Google repo
uni = unified_from_pretrained("mlnomad/gemma4-jax", dtype=jnp.bfloat16)
logits = uni.logits(input_ids)                                  # text
h = uni(input_ids, pixel_values=patches, image_position_ids=pos)  # + vision
h = uni(input_ids, input_features=frames)                       # + audio

From the original Google checkpoint

import jax.numpy as jnp
from gemma4_jax.convert import unified_from_safetensors
# accept the Gemma terms, then point at the official model.safetensors:
uni = unified_from_safetensors("path/to/google/gemma-4-12B/model.safetensors")

Re-serialize the port to a self-contained checkpoint yourself

from gemma4_jax.serialize import to_jax_safetensors
# streams the original -> NNX-keyed, sharded, bf16 safetensors (memory-safe)
to_jax_safetensors("model.safetensors", "out_dir/", shard_size_gb=5)

Tiny random model + smoke tests

import jax.numpy as jnp
from flax import nnx
import gemma4_jax as g4
model = g4.Gemma4TextModel(g4.tiny_config(), rngs=nnx.Rngs(0), dtype=jnp.float32)
logits = model.logits(jnp.zeros((1, 8), jnp.int32))

python -m gemma4_jax.smoke_test      # text
python -m gemma4_jax.smoke_test_mm   # vision + audio + exact param count

Multimodal (encoder-free) — implemented

Gemma 4 Unified has no vision tower and no audio tower. Raw inputs are projected straight into the 3840-dim decoder space and spliced into the text stream at placeholder token positions. All of this is now ported (multimodal.py

• unified.py), with tensor names verified against the real checkpoint:

• Vision (vision_embedder.* + embed_vision.*): raw 48×48×3 = 6912-pixel merged patches → LayerNorm → Dense(6912→3840) → LayerNorm → + factorized 2D positional embedding [1120,2,3840] (per-axis lookup, padding-masked) → LayerNorm → scale-free RMSNorm → Linear(3840→3840).
• Audio (embed_audio.*): raw 640-sample (40 ms @ 16 kHz) waveform frames → scale-free RMSNorm(640) → Linear(640→3840). No mel / conformer pipeline.
• Splice + mask: image/audio soft tokens are masked_scatter-ed into image_token_id/audio_token_id placeholder positions, and image/audio blocks attend bidirectionally (use_bidirectional_attention="vision").

import jax.numpy as jnp
from gemma4_jax.convert import unified_from_safetensors
from gemma4_jax.config import IMAGE_TOKEN_ID

uni = unified_from_safetensors("gemma4_jax/weights/model.safetensors")
# pixel_values: [B, num_patches, 6912]; image_position_ids: [B, num_patches, 2]
soft = uni.get_image_features(pixel_values, image_position_ids)   # [B, P, 3840]
# full forward with an [<txt> <IMG>*P <txt>] sequence:
h = uni(input_ids, pixel_values=pixel_values, image_position_ids=image_position_ids)

The image processor (patchify + 3×3 pool + position-id assignment) and audio feature extractor are not ported — feed pre-patchified pixel_values / pre-framed input_features, exactly as the HF processors emit them.

Files

• config.py — Gemma4TextConfig (real 12B) + vision/audio configs + token ids + tiny_config().
• rope.py — default + proportional/partial RoPE tables, split-half rotation.
• model.py — RMSNorm, MLP, dual attention, decoder layer, text model, param_count.
• multimodal.py — LayerNorm, multimodal embedder, vision patch embedder, vision/audio embedders.
• unified.py — Gemma4UnifiedModel (text + vision + audio splice + bidi mask), full param_count.
• convert.py — safetensors → NNX (load_into, load_multimodal_into, unified_from_safetensors, from_pretrained).
• generate.py — reference sampling loop (no KV cache).
• embed.py — use the backbone as a sentence-embedding model (last/mean pooling).
• smoke_test.py — text forward / causality / softcap / param-count checks.
• smoke_test_mm.py — vision/audio forward, splice, and exact full param-count check.