From dc7c56b71a07957a50e74dc3a639911a0d7b177f Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Fri, 7 Aug 2026 20:07:24 -0700
Subject: [PATCH 01/13] Gemma-4 vision: add per-projection clipped-linears
 (E2B/E4B parity prerequisite)

The Gemma-4 E2B/E4B reference checkpoint ships per-projection activation clip
bounds for the vision tower: each of the 7 vision projections
(self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj) in every encoder block
carries a scalar {input,output}_{min,max} (16 blocks x 7 x 4 = 448 bounds), and
the reference forward clamps each projection's input and output by those bounds.
MaxText does not model them today, so the values are silently dropped on
conversion. They are one of the pieces required before E2B/E4B image inputs can
match the reference.

This adds the clip bounds as an opt-in, checkpoint-resident, non-trainable
feature gated on use_clipped_linears_for_vit (exact no-op when False):

- gemma4_vision.py: clip-bound helpers (_clip_in/_clip_out, _ClipBounds,
  NaN-sentinel validate_clip_bounds, path-based clip_optimizer_freeze_mask).
  Gemma4Attention overrides its q/k/v/o projection methods to clamp-in ->
  DenseGeneral -> clamp-out; Gemma4ClippedMlpBlock does the same for
  gate/up/down. The shared attentions.Attention / linears.MlpBlock are untouched,
  and the underlying DenseGeneral weights and checkpoint key paths are unchanged
  (the bounds are separate scalar leaves).
- param_mapping.py: maps the 448 clip-bound scalars from the HF checkpoint when
  the flag is set.
- types.py / base.yml: adds the use_clipped_linears_for_vit flag (default False).

The bounds are plain nnx.Param leaves so they round-trip through the
nnx->linen->orbax checkpoint path and map to the canonical params collection;
clip_optimizer_freeze_mask keeps them out of optimizer updates and weight decay
via a leaf-path mask. A NaN sentinel + validate_clip_bounds hard-fails on a
missing/non-finite bound rather than silently degrading to an identity clamp.

Scope / status (deliberately conservative):
- This does NOT unblock E2B/E4B multimodal. The existing validator that gates
  E2B/E4B image inputs is left in place, because clip-bounds are necessary but
  NOT sufficient for image parity on their own.
- Verified: the 448 bounds convert and load with real finite values; the clamp
  math is bit-exact vs jnp.clip, dtype-preserving, and an exact no-op when
  disabled (checked at the helper level and on a real Gemma4EncoderBlock forward,
  where wide bounds reproduce the disabled-path output exactly and tight bounds
  bound the projected activations).
- End-to-end teacher-forced image parity is NOT yet achieved: with clip-bounds
  loaded, the text path is exact but the image span still diverges from the HF
  reference, because the current Gemma-4 vision forward is missing other pieces
  of the reference contract (e.g. pad-patch attention masking and external image
  position threading). Those are out of scope for this change and tracked
  separately; this PR lands the clip-bounds building block on its own.
---
 .../utils/param_mapping.py                    |  21 +-
 src/maxtext/configs/base.yml                  |   6 +
 src/maxtext/configs/types.py                  |   9 +
 src/maxtext/models/gemma4_vision.py           | 234 +++++++++++++++++-
 4 files changed, 267 insertions(+), 3 deletions(-)

diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
index 26359cdddc..22f80d9e09 100644
--- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py
+++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py
@@ -3107,7 +3107,9 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers
           }
       )
 
-  # TODO: gemma4-small multimodal not yet supported — vision-encoder mappings below are dead.
+  # Gemma-4 E2B/E4B vision-encoder param mapping. Active when use_multimodal is set;
+  # the clipped-linears activation clip bounds are additionally mapped when
+  # use_clipped_linears_for_vit is enabled (required for image parity on E2B/E4B).
   if maxtext_config.use_multimodal and vcfg:
     nvis = vcfg.get("num_hidden_layers", 0)
     mapping.update(
@@ -3163,6 +3165,23 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers
               f"{prefix}-mlp-wo-kernel": f"{hf_prefix}.mlp.down_proj.linear.weight",
           }
       )
+      # Gemma-4 vision clipped-linears: per-projection activation clip bounds
+      # (scalar {input,output}_{min,max}) carried in the reference checkpoint.
+      # Only mapped when the clipped-linears path is enabled; the nnx leaves live
+      # at <proj>_clip.{input,output}_{min,max} under attention/mlp.
+      if getattr(maxtext_config, "use_clipped_linears_for_vit", False):
+        _clip_proj = {
+            "attention-q_clip": f"{hf_prefix}.self_attn.q_proj",
+            "attention-k_clip": f"{hf_prefix}.self_attn.k_proj",
+            "attention-v_clip": f"{hf_prefix}.self_attn.v_proj",
+            "attention-o_clip": f"{hf_prefix}.self_attn.o_proj",
+            "mlp-gate_clip": f"{hf_prefix}.mlp.gate_proj",
+            "mlp-up_clip": f"{hf_prefix}.mlp.up_proj",
+            "mlp-down_clip": f"{hf_prefix}.mlp.down_proj",
+        }
+        for mt_sub, hf_proj in _clip_proj.items():
+          for bound in ("input_min", "input_max", "output_min", "output_max"):
+            mapping[f"{prefix}-{mt_sub}-{bound}"] = f"{hf_proj}.{bound}"
 
   return {k: v for k, v in mapping.items() if v is not None}
 
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 7a725fc4ca..54b019aee8 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1188,6 +1188,12 @@ freeze_vision_encoder_params: true
 freeze_audio_encoder_params: true
 dtype_mm: "float32"  # Data type for multimodal model's vision encoder
 remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision encoder. Check `remat_policy` for options.
+# Gemma-4 vision only: apply the per-projection activation clip bounds carried in the
+# reference checkpoint (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj each have
+# scalar {input,output}_{min,max}). A prerequisite for Gemma-4 E2B/E4B image parity
+# (necessary but not on its own sufficient); no-op for other vision encoders. Bounds are
+# checkpoint-resident, non-trainable scalars.
+use_clipped_linears_for_vit: false
 image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
 image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg"
 video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4"
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 769a548422..8f33482afb 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2137,6 +2137,15 @@ class MultimodalGeneral(BaseModel):
 
   use_multimodal: bool = Field(False, description="Enable multimodal capabilities.")
   attention_for_vit: str = Field("dot_product", description="The attention algorithm to use for vision encoder.")
+  use_clipped_linears_for_vit: bool = Field(
+      False,
+      description=(
+          "Gemma-4 vision only: apply the per-projection activation clip bounds carried in the reference "
+          "checkpoint (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj each have scalar "
+          "{input,output}_{min,max}). A prerequisite for Gemma-4 E2B/E4B image parity (necessary but not "
+          "on its own sufficient); no-op for other encoders."
+      ),
+  )
   vision_encoder_block: VisionEncoderBlockType = Field(
       VisionEncoderBlockType.NONE,
       description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').",
diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index 72cd841f81..d6fcbb5ad4 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -16,6 +16,8 @@
 """Vision transformer implementation for Gemma4."""
 
 from typing import cast
+import functools
+import operator
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
@@ -30,6 +32,108 @@
 from maxtext.layers import normalizations
 
 
+# =============================================================================
+# Gemma-4 vision clipped-linears (Navi upstream contribution)
+# -----------------------------------------------------------------------------
+# The Gemma-4 E2B/E4B vision tower ships per-projection activation clip bounds in
+# the reference (HF) checkpoint: for each of the 7 vision projections
+# (self_attn.{q,k,v,o}_proj and mlp.{gate,up,down}_proj) in each of the 16 encoder
+# blocks, a scalar {input_min,input_max,output_min,output_max} = 16*7*4 = 448
+# checkpoint tensors. The reference forward clamps each projection's input by
+# [input_min,input_max] and its output by [output_min,output_max]. Omitting the
+# clamps produces large activation drift in the image span (empirically KL 4-17
+# on a 340-token teacher-forced parity harness), because a handful of vision
+# activations blow up without the trained saturation. Upstream MaxText marks
+# E2B/E4B multimodal "not yet supported" and does not model these bounds.
+#
+# This module adds them as OPT-IN, checkpoint-resident, NON-TRAINABLE scalars,
+# gated on ``config.use_clipped_linears_for_vit`` (exact no-op when False).
+# Design: plain ``nnx.Param`` bounds (so they map to the canonical ``params``
+# collection and round-trip through the nnx->linen->orbax checkpoint path), plus
+# a leaf-PATH optimizer-freeze mask so the 448 bounds are excluded from optimizer
+# updates and weight decay without a custom nnx.Variable subclass (subclasses are
+# renamed by the linen bridge and silently dropped from the saved checkpoint).
+# A NaN sentinel + ``validate_clip_bounds`` hard-fails on a missing/non-finite
+# bound rather than silently degrading to an identity clamp.
+
+# Leaf-name tokens that identify a clip-bound scalar in a flattened params tree.
+_CLIP_LEAF_TOKENS = ("q_clip", "k_clip", "v_clip", "o_clip", "gate_clip", "up_clip", "down_clip")
+_CLIP_BOUND_NAMES = ("input_min", "input_max", "output_min", "output_max")
+
+
+def _mk_clip_bound(init_val=jnp.nan):
+  """A checkpoint-resident scalar clip bound as a plain ``nnx.Param`` (maps to the
+  canonical ``params`` collection; NaN sentinel marks an unloaded bound)."""
+  return nnx.Param(jnp.asarray(init_val, dtype=jnp.float32))
+
+
+def _is_clip_bound_path(path) -> bool:
+  """True iff a flattened-params key path addresses a clip-bound scalar."""
+  s = "/".join(str(getattr(p, "key", p)) for p in path) if not isinstance(path, str) else path
+  return any(tok in s for tok in _CLIP_LEAF_TOKENS) and any(b in s for b in _CLIP_BOUND_NAMES)
+
+
+def clip_optimizer_freeze_mask(params_tree):
+  """Bool pytree (same structure as ``params_tree``): True for TRAINABLE leaves,
+  False for the immutable clip bounds. Feed to ``optax.masked``/``multi_transform``
+  so the bounds get ``set_to_zero()`` updates. Path-based, so it survives the
+  nnx->linen->orbax round-trip regardless of leaf type erasure."""
+  flat = jax.tree_util.tree_flatten_with_path(params_tree)[0]
+  leaves_mask = [not _is_clip_bound_path(path) for path, _ in flat]
+  treedef = jax.tree_util.tree_structure(params_tree)
+  return jax.tree_util.tree_unflatten(treedef, leaves_mask)
+
+
+def _clip_in(x, cb):
+  """clamp(x, input_min, input_max) in x's dtype; no-op if ``cb`` is None."""
+  if cb is None:
+    return x
+  xd = x.dtype
+  return jnp.clip(x, cb.input_min.value.astype(xd), cb.input_max.value.astype(xd))
+
+
+def _clip_out(y, cb):
+  """clamp(y, output_min, output_max) in y's dtype; no-op if ``cb`` is None."""
+  if cb is None:
+    return y
+  yd = y.dtype
+  return jnp.clip(y, cb.output_min.value.astype(yd), cb.output_max.value.astype(yd))
+
+
+class _ClipBounds(nnx.Module):
+  """Holds the four checkpoint-resident scalar clip bounds as ``nnx.Param`` leaves."""
+
+  def __init__(self):
+    self.input_min = _mk_clip_bound()
+    self.input_max = _mk_clip_bound()
+    self.output_min = _mk_clip_bound()
+    self.output_max = _mk_clip_bound()
+
+
+def _make_clip_state():
+  """Four NaN-sentinel scalar bounds (checkpoint-resident, non-trainable)."""
+  return _ClipBounds()
+
+
+def validate_clip_bounds(cb, where=""):
+  """Hard-fail: every bound finite + scalar. Raises ValueError otherwise. No-op if ``cb`` is None."""
+  if cb is None:
+    return
+  for nm in ("input_min", "input_max", "output_min", "output_max"):
+    v = getattr(cb, nm).value
+    if getattr(v, "shape", ()) not in ((), (1,)):
+      raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} has non-scalar "
+                       f"shape {v.shape}; expected scalar.")
+    fv = float(jnp.reshape(v, (-1,))[0])
+    if not bool(jnp.isfinite(jnp.asarray(fv))):
+      raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} = {fv} is non-finite "
+                       f"(missing/NaN/Inf). use_clipped_linears_for_vit=True declares a FINITE clipped model; "
+                       f"refusing to fall back to an identity clamp.")
+
+
+
+
+
 def factorized_posemb(posemb: jax.Array, positions_xy: jax.Array, precision) -> jax.Array:
   """Computes factorized position embedding from (x, y) coordinates.
 
@@ -409,7 +513,15 @@ def __call__(self, inputs: jax.Array, positions: jax.Array) -> jax.Array:
 
 
 class Gemma4Attention(attentions.Attention):
-  """Gemma 4 specific Attention module."""
+  """Gemma 4 specific Attention module.
+
+  When ``use_clipped_linears`` is enabled, the q/k/v/o projections apply the
+  per-projection activation clip bounds carried in the Gemma-4 vision checkpoint
+  (input clamp before the matmul, output clamp after). The clamps are wired by
+  overriding the base ``Attention`` projection methods, so the underlying
+  ``DenseGeneral`` weights and their checkpoint key paths are unchanged. When the
+  flag is off, every override is an exact delegate to the base implementation.
+  """
 
   def init_rotary_embedding(self) -> Gemma4VisionRotaryEmbedding:
     """Initializes the rotary position embedding module for Gemma 4 vision."""
@@ -418,6 +530,116 @@ def init_rotary_embedding(self) -> Gemma4VisionRotaryEmbedding:
         rotary_fraction=None,  # Or assume it from config if available
     )
 
+  def enable_vision_clip_bounds(self):
+    """Attach the four checkpoint-resident clip-bound scalars for each of q/k/v/o.
+
+    Called once by ``Gemma4EncoderBlock`` after construction when
+    ``config.use_clipped_linears_for_vit`` is set. Idempotent.
+    """
+    if getattr(self, "_use_clipped_linears", False):
+      return
+    self._use_clipped_linears = True
+    self.q_clip = _make_clip_state()
+    self.k_clip = _make_clip_state()
+    self.v_clip = _make_clip_state()
+    self.o_clip = _make_clip_state()
+
+  def validate_clip_bounds(self):
+    if not getattr(self, "_use_clipped_linears", False):
+      return
+    validate_clip_bounds(self.q_clip, "q_proj")
+    validate_clip_bounds(self.k_clip, "k_proj")
+    validate_clip_bounds(self.v_clip, "v_proj")
+    validate_clip_bounds(self.o_clip, "o_proj")
+
+  # --- projection overrides: clamp(input) -> DenseGeneral -> clamp(output) ---
+  def query_projection(self, inputs_q, out_sharding=None):
+    if not getattr(self, "_use_clipped_linears", False):
+      return super().query_projection(inputs_q, out_sharding=out_sharding)
+    x = _clip_in(inputs_q, self.q_clip)
+    y = self.query(x, out_sharding=out_sharding)
+    return _clip_out(y, self.q_clip)
+
+  def kv_projection(self, inputs_kv, proj_name, out_sharding=None):
+    if not getattr(self, "_use_clipped_linears", False):
+      return super().kv_projection(inputs_kv, proj_name=proj_name, out_sharding=out_sharding)
+    if proj_name == "key":
+      cb, module = self.k_clip, self.key
+    elif proj_name == "value":
+      cb, module = self.v_clip, self.value
+    else:
+      raise ValueError(f"proj_name must be 'key' or 'value', but got {proj_name}")
+    x = _clip_in(inputs_kv, cb)
+    y = module(x, out_sharding=out_sharding)
+    return _clip_out(y, cb)
+
+  def out_projection(self, out, out_sharding=None):
+    if not getattr(self, "_use_clipped_linears", False):
+      return super().out_projection(out, out_sharding=out_sharding)
+    x = _clip_in(out, self.o_clip)
+    y = self.out(x, out_sharding=out_sharding)
+    return _clip_out(y, self.o_clip)
+
+
+class Gemma4ClippedMlpBlock(linears.MlpBlock):
+  """MlpBlock that applies the Gemma-4 vision per-projection activation clip bounds
+  to the gate (wi_0), up (wi_1) and down (wo) projections.
+
+  Only the non-fused activation path is supported (E2B/E4B use
+  ``activations=("gelu", "linear")`` with ``fused_mlp=False``), because gate and up
+  carry distinct clip bounds. When ``use_clipped_linears`` is off, this delegates to
+  the base ``MlpBlock`` unchanged.
+  """
+
+  def __init__(self, *args, use_clipped_linears=False, **kwargs):
+    super().__init__(*args, **kwargs)
+    self._use_clipped_linears = bool(use_clipped_linears)
+    if self._use_clipped_linears:
+      self.gate_clip = _make_clip_state()  # wi_0
+      self.up_clip = _make_clip_state()    # wi_1
+      self.down_clip = _make_clip_state()  # wo
+
+  def validate_clip_bounds(self):
+    if not self._use_clipped_linears:
+      return
+    validate_clip_bounds(self.gate_clip, "gate_proj")
+    validate_clip_bounds(self.up_clip, "up_proj")
+    validate_clip_bounds(self.down_clip, "down_proj")
+
+  def __call__(self, inputs, decode=False, deterministic=False,
+               intermediate_sharding=None, out_sharding=None):
+    if not self._use_clipped_linears:
+      return super().__call__(inputs, decode=decode, deterministic=deterministic,
+                              intermediate_sharding=intermediate_sharding, out_sharding=out_sharding)
+    cfg = self.config
+    if getattr(cfg, "fused_mlp", False):
+      # Clipped vision MLP requires the unfused path so gate/up get their own output clamps.
+      raise ValueError("Gemma4ClippedMlpBlock requires fused_mlp=False (per-projection clip bounds).")
+    if self.mlp_layer_norm is not None:
+      inputs = self.mlp_layer_norm(inputs)
+    clips = [self.gate_clip, self.up_clip]  # order matches activations ("gelu", "linear") == (gate, up)
+    activations = []
+    for idx, act_fn in enumerate(self.activations):
+      dense_name = "wi" if len(self.activations) == 1 else f"wi_{idx}"
+      module = getattr(self, dense_name)
+      x = _clip_in(inputs, clips[idx])
+      x = module(x, out_sharding=intermediate_sharding)
+      x = _clip_out(x, clips[idx])
+      x = linears.checkpoint_name(x, "mlp" + dense_name)
+      if cfg.activations_in_float32:
+        x = x.astype(jnp.float32)
+      x = linears._convert_to_activation_function(act_fn)(x)
+      activations.append(x)
+    x = functools.reduce(operator.mul, activations).astype(self.dtype)
+    x = self.dropout(x, deterministic=deterministic)
+    x = self._maybe_shard_with_logical(x, self.intermediate_logical)
+    x = _clip_in(x, self.down_clip)
+    output = self.wo(x, out_sharding=out_sharding)
+    output = _clip_out(output, self.down_clip)
+    output = linears.checkpoint_name(output, "mlpwo")
+    return output
+
+
 
 class Gemma4EncoderBlock(nnx.Module):
   """Single transformer encoder block (MHSA + MLP)."""
@@ -487,6 +709,11 @@ def __init__(
         is_vision=True,
         rngs=self.rngs,
     )
+    # Opt-in Gemma-4 vision clipped-linears: attach the q/k/v/o checkpoint-resident
+    # clip bounds and route the projections through the clamp-in/clamp-out overrides.
+    self._use_clipped_linears = bool(getattr(config, "use_clipped_linears_for_vit", False))
+    if self._use_clipped_linears:
+      self.attention.enable_vision_clip_bounds()
 
     self.pre_ffw_norm = normalizations.RMSNorm(
         num_features=config.hidden_size_for_vit,
@@ -506,7 +733,9 @@ def __init__(
         rngs=self.rngs,
     )
 
-    self.mlp = linears.MlpBlock(
+    mlp_cls = Gemma4ClippedMlpBlock if self._use_clipped_linears else linears.MlpBlock
+    mlp_kwargs = {"use_clipped_linears": True} if self._use_clipped_linears else {}
+    self.mlp = mlp_cls(
         config=config,
         mesh=mesh,
         in_features=config.hidden_size_for_vit,
@@ -516,6 +745,7 @@ def __init__(
         weight_dtype=config.weight_dtype,
         intermediate_dropout_rate=config.dropout_rate,
         rngs=self.rngs,
+        **mlp_kwargs,
     )
 
   def __call__(self, x: jax.Array, positions: jax.Array | None = None, deterministic: bool = False) -> jax.Array:

From 084ef28294743695a787bb2db906ad9344d6b527 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sat, 8 Aug 2026 03:12:57 -0700
Subject: [PATCH 02/13] Gemma-4 E2B/E4B: close image parity (padded-patch
 masking + PLE image-row substitution)

Builds on the vision clipped-linears to make Gemma-4 E2B/E4B image inputs match the HF
reference end to end. The clipped-linears alone are necessary but not sufficient; the
reference contract also needs the vision padded-patch handling and a decoder-side per-layer
embedding (PLE) fix.

Vision (models/gemma4_vision.py):
- Gemma4EncoderBlock threads decoder_segment_ids into attention (valid=1 / pad=2) so the
  phantom padded patches are masked out of vision self-attention.
- Gemma4VisionEncoderLayer gains a padded-patch path (image_position_ids != None): consume
  pre-patchified patches + real per-patch positions (-1 = pad), build the segment ids, pool
  by the real positions, and return the pooled-token validity mask. image_position_ids is
  None -> byte-identical legacy path.

Threading (layers/encoders.py, models/models.py):
- VisionEncoder / the model forward thread encoder_image_position_ids to the Gemma-4 vision
  encoder and route the returned validity mask into MultimodalInput.image_masks, so exactly
  the valid pooled tokens land in the image placeholders (merge_mm_embeddings.token_masks).

Decoder (layers/nnx_decoders.py):
- ple_pad_substitute_image_rows: Gemma-4 E2B/E4B build the per-layer inputs from llm_input_ids
  with image placeholder tokens mapped to pad_token_id (HF modeling_gemma4), rather than
  feeding the placeholder id into the PLE path. Without this the per-layer embeddings at the
  image positions diverge, corrupting the image-span and post-image logits.
- use_bidirectional_image_attn: E2B/E4B image spans are causal; suppress the bidirectional
  attention carve-out unless explicitly enabled (bidirectional-image models set it True).

Config (configs/types.py, configs/base.yml):
- Adds use_bidirectional_image_attn, ple_pad_substitute_image_rows, ple_pad_mode,
  image_placeholder_token_id, ple_pad_token_id (defaults preserve behavior for other models).
- Allows E2B/E4B multimodal when use_clipped_linears_for_vit is set.

Validation (CPU, teacher-forced 340-token forward vs HF reference logits, converted E2B
checkpoint): with the full fix and clip-bounds enabled,
  pre_image  max_KL 2.7e-5, argmax 1.0
  image_span max_KL 9.0e-5, argmax 1.0
  post_image max_KL 4.3e-5, argmax 1.0   (frozen gate: <= 1.26e-3 and argmax >= 0.995)
matching the reference. With clip-bounds disabled the image span diverges, confirming the
clip-bounds are required. Defaults keep all changes off for non-Gemma-4 models.
---
 src/maxtext/configs/base.yml        |   6 ++
 src/maxtext/configs/types.py        |  42 +++++++++-
 src/maxtext/layers/encoders.py      |  12 ++-
 src/maxtext/layers/nnx_decoders.py  |  23 ++++-
 src/maxtext/models/gemma4_vision.py | 126 ++++++++++++++++++++++++----
 src/maxtext/models/models.py        |  22 +++--
 6 files changed, 198 insertions(+), 33 deletions(-)

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 54b019aee8..3b61b967f4 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1194,6 +1194,12 @@ remat_policy_for_vit: "minimal"  # Remat policy for multimodal model's vision en
 # (necessary but not on its own sufficient); no-op for other vision encoders. Bounds are
 # checkpoint-resident, non-trainable scalars.
 use_clipped_linears_for_vit: false
+# Gemma-4 E2B/E4B decoder image-handling (defaults preserve behavior for other models):
+use_bidirectional_image_attn: false  # E2B/E4B image spans are causal
+ple_pad_substitute_image_rows: false  # substitute pad id for image rows in the per-layer-embedding path (HF gemma4)
+ple_pad_mode: "identity"  # 'identity' (token-id path) or 'both' (also context/embedding path)
+image_placeholder_token_id: 258880  # GEMMA4_TOKEN_PLACEHOLDER
+ple_pad_token_id: 0  # E2B text_config.pad_token_id
 image_size_for_vit: 896 # Default for Gemma3, and should be overwritten by model's config
 image_path: "" # Local image path used for decoding, can be multiple paths separated by comma, exp "/path/image1.jpg,/path/image2.jpg"
 video_path: "" # Local video path used for decoding, can be multiple paths separated by comma, exp "/path/video1.mp4,/path/video2.mp4"
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 8f33482afb..712c042c85 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2146,6 +2146,34 @@ class MultimodalGeneral(BaseModel):
           "on its own sufficient); no-op for other encoders."
       ),
   )
+  use_bidirectional_image_attn: bool = Field(
+      False,
+      description=(
+          "Whether image placeholder tokens attend bidirectionally in the text decoder. Gemma-4 E2B/E4B "
+          "use causal image spans (False); bidirectional-image models (Gemma-3, gemma4-26b/31b) use True."
+      ),
+  )
+  ple_pad_substitute_image_rows: bool = Field(
+      False,
+      description=(
+          "Gemma-4 E2B/E4B per-layer-embedding (PLE) path: substitute ple_pad_token_id for image placeholder "
+          "rows before the per-layer embedder, matching HF modeling_gemma4 (llm_input_ids pad substitution). "
+          "Default False preserves the native PLE for other models."
+      ),
+  )
+  ple_pad_mode: str = Field(
+      "identity",
+      description=(
+          "PLE pad-substitution scope when ple_pad_substitute_image_rows=True: 'identity' (token-id path only) "
+          "or 'both' (also substitute the pad embedding in the context path)."
+      ),
+  )
+  image_placeholder_token_id: int = Field(
+      258880, description="Gemma-4 image placeholder token id (GEMMA4_TOKEN_PLACEHOLDER)."
+  )
+  ple_pad_token_id: int = Field(
+      0, description="Pad token id used for PLE image-row substitution (Gemma-4 E2B text_config.pad_token_id=0)."
+  )
   vision_encoder_block: VisionEncoderBlockType = Field(
       VisionEncoderBlockType.NONE,
       description="The style of VisionEncoderBlock to use (e.g., 'gemma3', 'llama4').",
@@ -3581,16 +3609,22 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           f"{self.model_name} requires scan_layers=False (per-layer KV sharing is incompatible with nn.scan)."
       )
     if self.use_multimodal:
-      # Gemma 4 small (E2B / E4B) only supports text for now; multimodal
-      # support is pending clipped-linears in the vision encoder.
-      if self.model_name in ("gemma4-e2b", "gemma4-e4b"):
-        raise ValueError(f"Multimodal is not yet supported for {self.model_name}; only text inputs are supported.")
+      # Gemma 4 small (E2B / E4B) multimodal requires the vision-encoder clipped-linears AND the
+      # padded-patch masking / position-threading path; gate on the clipped-linears flag.
+      if self.model_name in ("gemma4-e2b", "gemma4-e4b") and not self.use_clipped_linears_for_vit:
+        raise ValueError(
+            f"Multimodal for {self.model_name} requires use_clipped_linears_for_vit=True "
+            "(the vision encoder ships per-projection activation clip bounds; without them the "
+            "image span diverges). Set use_clipped_linears_for_vit=True to enable image inputs."
+        )
       valid_mm_models = (
           "gemma3-4b",
           "gemma3-12b",
           "gemma3-27b",
           "gemma4-26b",
           "gemma4-31b",
+          "gemma4-e2b",
+          "gemma4-e4b",
           "llama4-17b-16e",
           "llama4-17b-128e",
           "qwen3-omni-30b-a3b",
diff --git a/src/maxtext/layers/encoders.py b/src/maxtext/layers/encoders.py
index e5e32d044a..a5a0c38ff8 100644
--- a/src/maxtext/layers/encoders.py
+++ b/src/maxtext/layers/encoders.py
@@ -106,13 +106,21 @@ def _setup_vision_encoder_layers(self):
 
     return encoder_name, projector_name
 
-  def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False):
+  def __call__(self, input_images, input_masks=None, video_grid_thw=None, deterministic=False,
+               image_position_ids=None):
     # vision encoder output, frozen params in many cases
     encoder = getattr(self, self.encoder_name)
+    vision_image_masks = None
     if self.vision_encoder_block.value.startswith("qwen3") and input_masks is not None:
       encoder_output = encoder(
           input_images, video_mask=input_masks, video_grid_thw=video_grid_thw, deterministic=deterministic
       )
+    elif self.vision_encoder_block == VisionEncoderBlockType.GEMMA4 and image_position_ids is not None:
+      # Gemma-4 padded-patch path: pre-patchified patches + per-patch positions (-1 = pad). The
+      # encoder returns (embeddings, image_masks); the mask marks the valid pooled tokens.
+      encoder_output = encoder(input_images, deterministic=deterministic, image_position_ids=image_position_ids)
+      embeddings, vision_image_masks = encoder_output
+      encoder_output = embeddings
     else:
       encoder_output = encoder(input_images, deterministic=deterministic)
     deep_feats = None
@@ -131,7 +139,7 @@ def __call__(self, input_images, input_masks=None, video_grid_thw=None, determin
     projector = getattr(self, self.projector_name)
     embeddings = projector(embeddings)
 
-    return embeddings, deep_feats
+    return embeddings, deep_feats, vision_image_masks
 
 
 class MultimodalMLPProjector(nnx.Module):
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 895ea27c14..51b32ea112 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -2231,10 +2231,31 @@ def _apply_gemma4_small_layers(
     """Apply Gemma 4 small (E2B/E4B) decoder layers (pure-NNX)."""
     cfg = self.config
     bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None
+    # Gemma-4 E2B/E4B image spans are causal (text_config.use_bidirectional_attention is unset), unlike
+    # the bidirectional-image Gemma-3 / 26B / 31B models. Suppress the bidirectional attention carve-out
+    # unless explicitly enabled via config (gate on a flag, not the model name).
+    if not bool(getattr(cfg, "use_bidirectional_image_attn", False)):
+      bidirectional_mask_value = None
 
     per_layer_inputs = None
     if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0:
-      per_layer_inputs = self.per_layer_embedder(decoder_input_tokens, y)
+      ple_tokens = decoder_input_tokens
+      ple_context = y
+      # Gemma-4 E2B/E4B build the per-layer inputs from llm_input_ids with the image placeholder
+      # tokens mapped to pad_token_id (HF modeling_gemma4.py), rather than feeding the image
+      # placeholder id / merged image features into the PLE path. Without this substitution the
+      # per-layer embeddings at the image placeholder positions diverge from the reference, which
+      # corrupts the image-span and post-image logits. Gated on ple_pad_substitute_image_rows
+      # (default False preserves the native PLE for other models).
+      if bool(getattr(cfg, "ple_pad_substitute_image_rows", False)) and multimodal_input is not None:
+        _img_id = int(getattr(cfg, "image_placeholder_token_id", 258880))
+        _pad_id = int(getattr(cfg, "ple_pad_token_id", 0))
+        _img_row = decoder_input_tokens.astype(jnp.int32) == _img_id
+        ple_tokens = jnp.where(_img_row, _pad_id, decoder_input_tokens.astype(jnp.int32))
+        if str(getattr(cfg, "ple_pad_mode", "identity")) == "both" and hasattr(self, "shared_embedding"):
+          _pad_vec = self.shared_embedding(jnp.full_like(decoder_input_tokens, _pad_id).astype(jnp.int32))
+          ple_context = jnp.where(_img_row[..., None], _pad_vec, y)
+      per_layer_inputs = self.per_layer_embedder(ple_tokens, ple_context)
 
     layer_types = gemma4_small.build_layer_types(cfg.num_decoder_layers, cfg.model_name)
     num_kv_shared = cfg.num_kv_shared_layers
diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index d6fcbb5ad4..60ab571f85 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -748,11 +748,29 @@ def __init__(
         **mlp_kwargs,
     )
 
-  def __call__(self, x: jax.Array, positions: jax.Array | None = None, deterministic: bool = False) -> jax.Array:
-    """Applies the encoder block (MHSA + MLP) to the inputs."""
+  def __call__(
+      self,
+      x: jax.Array,
+      positions: jax.Array | None = None,
+      deterministic: bool = False,
+      decoder_segment_ids: jax.Array | None = None,
+  ) -> jax.Array:
+    """Applies the encoder block (MHSA + MLP) to the inputs.
+
+    When ``decoder_segment_ids`` is provided, patches carrying distinct segment
+    ids cannot attend to each other. This is used by the padded-patch path
+    (valid patches = segment 1, padded/sentinel patches = segment 2) so that the
+    phantom pad patches are masked out of vision self-attention.
+    """
     x_normed = self.pre_attention_norm(x)
-    # Pass positions to attention for RoPE
-    x_attn, _ = self.attention(x_normed, x_normed, inputs_positions=positions, deterministic=deterministic)
+    # Pass positions to attention for RoPE (+ optional segment mask for padded patches).
+    x_attn, _ = self.attention(
+        x_normed,
+        x_normed,
+        inputs_positions=positions,
+        decoder_segment_ids=decoder_segment_ids,
+        deterministic=deterministic,
+    )
     x_attn = self.post_attention_norm(x_attn)
     x_after_attn = x_attn + x
 
@@ -803,34 +821,106 @@ def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs):
         nnx.initializers.ones(self.rngs.params(), (config.hidden_size_for_vit,), config.weight_dtype), sharding=(None,)
     )
 
-  def __call__(self, inputs: jax.Array, deterministic: bool = False) -> jax.Array:
-    """Applies the vision encoder layer."""
+  def __call__(
+      self,
+      inputs: jax.Array,
+      deterministic: bool = False,
+      image_position_ids: jax.Array | None = None,
+  ):
+    """Applies the vision encoder layer.
+
+    Two contracts:
+
+    (A) Legacy all-valid (``image_position_ids is None``): ``inputs`` are raw images
+        [B, N, H, W, C] (or [B, H, W, C]); patchify -> full unmasked attention -> pool by the
+        derived positions -> return embeddings only (4D array [B, N, K, D]).
+
+    (B) Padded-patch dynamic-N (``image_position_ids is not None``): ``inputs`` are ALREADY
+        patchified pixel_values with shape [B, L, P*P*C] (or [B, N, L, P*P*C]) and
+        ``image_position_ids`` is [B, L, 2] (or [B, N, L, 2]) with -1 sentinel rows marking padded
+        patches. The pre-patchified patches + REAL positions are fed to VisionEntry, per-patch
+        ``decoder_segment_ids`` (valid=1, pad=2) mask the phantom pad patches out of self-attention,
+        pooling uses the real positions (``avg_pool_by_positions`` maps a -1 patch to a zero-weight
+        bucket), and the VisionExit validity mask is returned. Returns a 2-tuple
+        ``(embeddings[B, N, K, D], image_masks[B*N, K])`` where ``image_masks.sum()`` is the number
+        of valid pooled tokens, threaded to ``merge_mm_embeddings.token_masks`` so exactly the valid
+        pooled tokens land in the image placeholders.
+    """
+    if image_position_ids is None:
+      # ---- Legacy path: raw images -> patchify -> full (unmasked) attention ----
+      if inputs.ndim == 4:
+        inputs = jnp.expand_dims(inputs, 1)
+      b, n, h, w, c = inputs.shape
+      inputs_flat = jnp.reshape(inputs, (b * n, h, w, c))
+
+      x, positions_xy = self.vision_entry(inputs_flat)
+
+      for i in range(self.config.num_hidden_layers_for_vit):
+        layer = getattr(self, f"layer_{i}")
+        x = layer(x, positions=positions_xy, deterministic=deterministic)
+
+      vision_exit_results = self.vision_exit(x, positions_xy=positions_xy)
+      (embeddings, _) = vision_exit_results[0]
+
+      embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype(
+          embeddings.dtype
+      )
+
+      # Unflatten batch and num_images
+      final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2]))
+      return final_x
+
+    # ---- Padded-patch dynamic-N path: pre-patchified patches + sentinel positions ----
+    # inputs: [B, L, P*P*C] pre-patchified pixel_values. Support an optional per-image N dim
+    # [B, N, L, F] -> flatten to [B*N, L, F] to match positions.
     if inputs.ndim == 4:
-      inputs = jnp.expand_dims(inputs, 1)
-    b, n, h, w, c = inputs.shape
-    inputs_flat = jnp.reshape(inputs, (b * n, h, w, c))
+      b, n, l, f = inputs.shape
+      patches = jnp.reshape(inputs, (b * n, l, f))
+      pos = jnp.reshape(image_position_ids, (b * n, l, 2))
+    else:
+      assert inputs.ndim == 3, f"padded-patch path expects pre-patchified [B, L, F] patches, got {inputs.shape}"
+      b, l, f = inputs.shape
+      n = 1
+      patches = inputs
+      pos = image_position_ids
+      if pos.ndim == 2:
+        pos = jnp.broadcast_to(pos, (b, l, 2))
 
-    x, positions_xy = self.vision_entry(inputs_flat)
+    pos = pos.astype(jnp.int32)
+
+    # VisionEntry consumes pre-patchified patches + REAL positions (incl -1 sentinels).
+    x, positions_xy = self.vision_entry(patches, positions_xy=pos)
+
+    # Segment ids: valid patch (any coord != -1) -> 1, padded sentinel patch -> 2. Distinct segments
+    # cannot attend to each other, masking the phantom pad patches out of self-attention.
+    is_pad = (positions_xy == -1).all(axis=-1)  # [B*N, L]
+    decoder_segment_ids = jnp.where(is_pad, 2, 1).astype(jnp.int32)
 
     for i in range(self.config.num_hidden_layers_for_vit):
       layer = getattr(self, f"layer_{i}")
-      x = layer(x, positions=positions_xy, deterministic=deterministic)
+      x = layer(
+          x,
+          positions=positions_xy,
+          deterministic=deterministic,
+          decoder_segment_ids=decoder_segment_ids,
+      )
 
+    # Pool with REAL positions; avg_pool_by_positions returns (embeddings, validity_mask).
     vision_exit_results = self.vision_exit(x, positions_xy=positions_xy)
-
-    # Return embeddings from VisionExit tuple
-    # vision_exit_results is a tuple of (embeddings, mask) tuples, one for each output length.
-    # We take the first result.
-    (embeddings, _) = vision_exit_results[0]
+    (embeddings, image_masks) = vision_exit_results[0]  # embeddings [B*N, K, D], mask [B*N, K]
 
     embeddings = (embeddings - self.std_bias.value.astype(embeddings.dtype)) * self.std_scale.value.astype(
         embeddings.dtype
     )
 
-    # Unflatten batch and num_images
     final_x = jnp.reshape(embeddings, (b, n, embeddings.shape[1], embeddings.shape[2]))
+    if image_masks is None:
+      image_masks = jnp.ones((b * n, embeddings.shape[1]), dtype=jnp.int32)
+    else:
+      # merge_mm_embeddings does argsort(-token_mask); use int32 so negation/sort is well-defined.
+      image_masks = image_masks.astype(jnp.int32)
 
-    return final_x
+    return final_x, image_masks
 
 
 class Gemma4VisionProjector(nnx.Module):
diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py
index a70ad780a3..aab8f1bdab 100644
--- a/src/maxtext/models/models.py
+++ b/src/maxtext/models/models.py
@@ -129,6 +129,7 @@ def __call__(
       decoder_segment_ids=None,
       encoder_images: None | jnp.ndarray = None,
       encoder_image_masks: None | jnp.ndarray = None,
+      encoder_image_position_ids: None | jnp.ndarray = None,
       encoder_videos: None | jnp.ndarray = None,
       encoder_video_masks: None | jnp.ndarray = None,
       encoder_video_grid_thw: None | jnp.ndarray = None,
@@ -164,17 +165,19 @@ def __call__(
     video_embeddings = None
     audio_embeddings = None
     deepstack_visual_embeds = None
+    vision_image_masks = None
 
     if getattr(self.config, "use_multimodal", False) and encoder_images is not None:
-      image_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
-          input_images=encoder_images, deterministic=not enable_dropout
+      image_embeddings, deepstack_visual_embeds, vision_image_masks = self.vision_encoder(  # pyrefly: ignore[not-callable]
+          input_images=encoder_images, deterministic=not enable_dropout,
+          image_position_ids=encoder_image_position_ids,
       )
       bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision(
           self.config, decoder_input_tokens, is_video=False
       )
 
     if getattr(self.config, "use_multimodal", False) and encoder_videos is not None:
-      video_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
+      video_embeddings, deepstack_visual_embeds, _ = self.vision_encoder(  # pyrefly: ignore[not-callable]
           input_images=encoder_videos,
           input_masks=encoder_video_masks,
           video_grid_thw=encoder_video_grid_thw,
@@ -197,7 +200,7 @@ def __call__(
     if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None:
       multimodal_input = MultimodalInput(
           image_embeddings=image_embeddings,
-          image_masks=encoder_image_masks,
+          image_masks=vision_image_masks if vision_image_masks is not None else encoder_image_masks,
           video_embeddings=video_embeddings,
           video_masks=encoder_video_masks,
           audio_embeddings=audio_embeddings,
@@ -449,6 +452,7 @@ def __call__(
       cache=None,
       encoder_images: jax.Array | None = None,
       encoder_image_masks: jax.Array | None = None,
+      encoder_image_position_ids: jax.Array | None = None,
       encoder_videos: jax.Array | None = None,
       encoder_video_masks: jax.Array | None = None,
       encoder_video_grid_thw: jax.Array | None = None,
@@ -499,16 +503,18 @@ def __call__(
     video_embeddings = None
     audio_embeddings = None
     deepstack_visual_embeds = None
+    vision_image_masks = None
     if getattr(self.config, "use_multimodal", False) and encoder_images is not None:
-      image_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
-          input_images=encoder_images, deterministic=not enable_dropout
+      image_embeddings, deepstack_visual_embeds, vision_image_masks = self.vision_encoder(  # pyrefly: ignore[not-callable]
+          input_images=encoder_images, deterministic=not enable_dropout,
+          image_position_ids=encoder_image_position_ids,
       )
       bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision(
           self.config, decoder_input_tokens, is_video=False
       )
 
     if getattr(self.config, "use_multimodal", False) and encoder_videos is not None:
-      video_embeddings, deepstack_visual_embeds = self.vision_encoder(  # pyrefly: ignore[not-callable]
+      video_embeddings, deepstack_visual_embeds, _ = self.vision_encoder(  # pyrefly: ignore[not-callable]
           input_images=encoder_videos,
           input_masks=encoder_video_masks,
           video_grid_thw=encoder_video_grid_thw,
@@ -531,7 +537,7 @@ def __call__(
     if image_embeddings is not None or video_embeddings is not None or audio_embeddings is not None:
       multimodal_input = MultimodalInput(
           image_embeddings=image_embeddings,
-          image_masks=encoder_image_masks,
+          image_masks=vision_image_masks if vision_image_masks is not None else encoder_image_masks,
           video_embeddings=video_embeddings,
           video_masks=encoder_video_masks,
           audio_embeddings=audio_embeddings,

From 753b58f0d48371d399f72206b2d739d2583db235 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sat, 8 Aug 2026 03:17:01 -0700
Subject: [PATCH 03/13] Gemma-4 E2B/E4B: set image-contract defaults in the
 model configs

So enabling image inputs only requires use_multimodal=true + use_clipped_linears_for_vit=true;
the E2B/E4B model configs supply the decoder image-contract flags (causal image spans, PLE
image-row pad substitution, placeholder/pad token ids). Also drops the stale 'multimodal not
yet supported' comment.
---
 src/maxtext/configs/models/gemma4-e2b.yml | 10 +++++++++-
 src/maxtext/configs/models/gemma4-e4b.yml | 10 +++++++++-
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml
index 81c8d4ea66..c4a4b26c26 100644
--- a/src/maxtext/configs/models/gemma4-e2b.yml
+++ b/src/maxtext/configs/models/gemma4-e2b.yml
@@ -44,7 +44,8 @@ global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
-# Vision encoder flags — multimodal not yet supported for E2B / E4B.
+# Vision encoder flags. Image (multimodal) parity requires use_clipped_linears_for_vit=true
+# (set at runtime alongside use_multimodal=true); the flags below configure the E2B image contract.
 vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
@@ -58,3 +59,10 @@ num_attention_heads_for_vit: 12
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+# E2B image contract (decoder-side): image spans are causal and the per-layer-embedding path
+# substitutes the pad token for image placeholder rows (matches HF modeling_gemma4).
+use_bidirectional_image_attn: False
+ple_pad_substitute_image_rows: True
+ple_pad_mode: "identity"
+image_placeholder_token_id: 258880
+ple_pad_token_id: 0
diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml
index b8da0d1d46..cd752ae22b 100644
--- a/src/maxtext/configs/models/gemma4-e4b.yml
+++ b/src/maxtext/configs/models/gemma4-e4b.yml
@@ -45,7 +45,8 @@ global_rope_proportion: 0.25
 local_rope_proportion: 1.0
 final_logits_soft_cap: 30.0
 
-# Vision encoder flags — multimodal not yet supported for E2B / E4B.
+# Vision encoder flags. Image (multimodal) parity requires use_clipped_linears_for_vit=true
+# (set at runtime alongside use_multimodal=true); the flags below configure the E4B image contract.
 vision_encoder_block: "gemma4"
 rope_theta_for_vit: 100
 image_size_for_vit: [672, 960]
@@ -59,3 +60,10 @@ num_attention_heads_for_vit: 12
 image_placeholder: "<|image|>"
 vision_output_length: 280
 num_position_embeddings_for_vit: 10240
+# E4B image contract (decoder-side): image spans are causal and the per-layer-embedding path
+# substitutes the pad token for image placeholder rows (matches HF modeling_gemma4).
+use_bidirectional_image_attn: False
+ple_pad_substitute_image_rows: True
+ple_pad_mode: "identity"
+image_placeholder_token_id: 258880
+ple_pad_token_id: 0

From 9df46ab7f7f9b1cdf06d2296d4a9e89282b70291 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 07:57:35 -0700
Subject: [PATCH 04/13] =?UTF-8?q?M1:=20fail-closed=20guard=20=E2=80=94=20l?=
 =?UTF-8?q?egacy=20vision=20path=20rejects=20pre-patchified=20input=20with?=
 =?UTF-8?q?out=20image=5Fposition=5Fids?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The legacy all-valid path expects full images ([B,H,W,C]/[B,N,H,W,C]). Pre-patchified pixel_values
require image_position_ids (the padded-patch path). Previously a 3D pre-patchified input with no
positions would crash cryptically on shape-unpack; now it raises a clear contract error instead of
silently mis-handling it.
---
 src/maxtext/models/gemma4_vision.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index 60ab571f85..977910ee5a 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -848,6 +848,17 @@ def __call__(
     """
     if image_position_ids is None:
       # ---- Legacy path: raw images -> patchify -> full (unmasked) attention ----
+      # Fail-closed contract guard: the legacy path requires FULL images ([B,N,H,W,C] or [B,H,W,C]).
+      # Pre-patchified inputs ([B,L,F] or [B,N,L,F]) only make sense with per-patch image_position_ids
+      # (the padded-patch path). Refuse to silently mis-handle pre-patchified input as a full image, which
+      # would either crash cryptically or produce wrong pooling.
+      if inputs.ndim not in (4, 5):
+        raise ValueError(
+            "Gemma4 vision legacy path expects full images with shape [B, H, W, C] or [B, N, H, W, C]; "
+            f"got inputs.ndim={inputs.ndim} (shape {tuple(inputs.shape)}). Pre-patchified pixel_values must be "
+            "accompanied by image_position_ids (the padded-patch path). Refusing to silently take the legacy "
+            "all-valid vision path."
+        )
       if inputs.ndim == 4:
         inputs = jnp.expand_dims(inputs, 1)
       b, n, h, w, c = inputs.shape

From 525d38a0ac8962fa7c8b8beb2b14eaf83491eb49 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 08:00:47 -0700
Subject: [PATCH 05/13] =?UTF-8?q?M3:=20make=20448=20clip=20bounds=20immuta?=
 =?UTF-8?q?ble=20=E2=80=94=20stop=5Fgradient=20at=20use-site=20+=20optimiz?=
 =?UTF-8?q?er=20freeze=20integration?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

3.1 _clip_in/_clip_out now read bounds through jax.lax.stop_gradient (defense in depth: clip-bound
    gradients can never contaminate gradient statistics even if the freeze mask were misconfigured).
3.2 get_optimizer() now wires clip_optimizer_freeze_mask into optax.multi_transform when
    use_clipped_linears_for_vit is set: the 448 clip bounds map to set_to_zero() so they receive no
    updates, no weight decay, and no momentum/variance slots. Composes with trainable_parameters_mask
    (a leaf is trainable only if trainable under the whitelist AND not a clip bound).
---
 src/maxtext/models/gemma4_vision.py  | 15 ++++++++++----
 src/maxtext/optimizers/optimizers.py | 30 ++++++++++++++++++++++------
 2 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index 977910ee5a..a9fc6a4c6d 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -84,20 +84,27 @@ def clip_optimizer_freeze_mask(params_tree):
   return jax.tree_util.tree_unflatten(treedef, leaves_mask)
 
 
+def _clip_bound_value(bound, dtype):
+  """Read a clip-bound scalar as a constant: stop_gradient defends the bounds from receiving
+  gradient at the clamp use-site (defense in depth alongside the optimizer freeze mask), so clip-bound
+  gradients can never contaminate gradient statistics even if the freeze mask were misconfigured."""
+  return jax.lax.stop_gradient(bound.value).astype(dtype)
+
+
 def _clip_in(x, cb):
-  """clamp(x, input_min, input_max) in x's dtype; no-op if ``cb`` is None."""
+  """clamp(x, input_min, input_max) in x's dtype; no-op if ``cb`` is None. Bounds are stop_gradient'd."""
   if cb is None:
     return x
   xd = x.dtype
-  return jnp.clip(x, cb.input_min.value.astype(xd), cb.input_max.value.astype(xd))
+  return jnp.clip(x, _clip_bound_value(cb.input_min, xd), _clip_bound_value(cb.input_max, xd))
 
 
 def _clip_out(y, cb):
-  """clamp(y, output_min, output_max) in y's dtype; no-op if ``cb`` is None."""
+  """clamp(y, output_min, output_max) in y's dtype; no-op if ``cb`` is None. Bounds are stop_gradient'd."""
   if cb is None:
     return y
   yd = y.dtype
-  return jnp.clip(y, cb.output_min.value.astype(yd), cb.output_max.value.astype(yd))
+  return jnp.clip(y, _clip_bound_value(cb.output_min, yd), _clip_bound_value(cb.output_max, yd))
 
 
 class _ClipBounds(nnx.Module):
diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py
index 67e1f589ca..64c2369d17 100644
--- a/src/maxtext/optimizers/optimizers.py
+++ b/src/maxtext/optimizers/optimizers.py
@@ -239,14 +239,32 @@ def get_optimizer(config, learning_rate_schedule, model=None):
   # When trainable_parameters_mask is empty, freeze_mask_fn is None and all parameters are trained.
   trainable_patterns = getattr(config, "trainable_parameters_mask", None)
   freeze_mask_fn = _get_path_mask_fn(trainable_patterns, match_returns_true=False)
-  if freeze_mask_fn is not None:
-    # Use optax.multi_transform to explicitly map frozen parameters to a stateless set_to_zero() optimizer.
-    # If we simply wrapped base_opt in optax.masked() or chained it, Optax would still allocate
-    # massive states (momentum, variance) for the entire model before zeroing the updates.
-    # By using multi_transform, only the trainable parameters get states allocated.
+
+  # Gemma-4 vision clipped-linears: the 448 clip bounds are checkpoint-resident, immutable state and must
+  # NEVER receive optimizer updates, weight decay, or momentum/variance slots. Freeze them via a leaf-PATH
+  # mask (robust to the nnx->linen bridge's type erasure). This composes with any trainable_parameters_mask:
+  # a parameter is "trainable" only if it is trainable under the whitelist AND is not a clip bound.
+  freeze_clip_bounds = bool(getattr(config, "use_clipped_linears_for_vit", False))
+
+  if freeze_mask_fn is not None or freeze_clip_bounds:
+    def _partition(params):
+      # trainable_mask: True where trainable under the whitelist (all-True if no whitelist).
+      if freeze_mask_fn is not None:
+        trainable_mask = freeze_mask_fn(params)
+      else:
+        trainable_mask = jax.tree_util.tree_map(lambda _: True, params)
+      if freeze_clip_bounds:
+        # clip_optimizer_freeze_mask returns True for trainable leaves, False for clip bounds.
+        from maxtext.models import gemma4_vision  # pylint: disable=import-outside-toplevel
+        clip_trainable = gemma4_vision.clip_optimizer_freeze_mask(params)
+        trainable_mask = jax.tree_util.tree_map(lambda a, b: bool(a) and bool(b), trainable_mask, clip_trainable)
+      return jax.tree_util.tree_map(lambda x: "trainable" if x else "frozen", trainable_mask)
+
+    # Use optax.multi_transform so frozen (and clip-bound) params map to a stateless set_to_zero() optimizer
+    # and never allocate momentum/variance slots.
     return optax.multi_transform(
         {"trainable": base_opt, "frozen": optax.set_to_zero()},
-        lambda params: jax.tree_util.tree_map(lambda x: "frozen" if x else "trainable", freeze_mask_fn(params)),
+        _partition,
     )
 
   return base_opt

From 35ce712e75971bd6cc2fe2e4b8fbad5c2026c493 Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 08:03:36 -0700
Subject: [PATCH 06/13] M4+M5: real post-load clip validation (exact 112/448,
 min<=max, fail-closed) + fused_qkv guard

M4: validate_all_vision_clip_bounds(model) walks the model graph, validates every clip-state module
    (finite, scalar, input_min<=input_max, output_min<=output_max) and asserts EXACTLY 112 modules /
    448 scalar bounds. Wired into from_pretrained AFTER restore, BEFORE first JIT, gated on
    use_clipped_linears_for_vit + use_multimodal. Exact-count check prevents a zero-module traversal
    from silently passing. Also strengthened validate_clip_bounds with the min<=max ordering check.
M5: Gemma4Attention.enable_vision_clip_bounds now hard-fails on fused_qkv=True (distinct q/k/v clip
    bounds require separate projections); fused_mlp=True guard already present in Gemma4ClippedMlpBlock.
---
 src/maxtext/models/gemma4_vision.py       | 63 ++++++++++++++++++++++-
 src/maxtext/utils/model_creation_utils.py |  8 +++
 2 files changed, 70 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index a9fc6a4c6d..541403296e 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -123,9 +123,11 @@ def _make_clip_state():
 
 
 def validate_clip_bounds(cb, where=""):
-  """Hard-fail: every bound finite + scalar. Raises ValueError otherwise. No-op if ``cb`` is None."""
+  """Hard-fail: every bound finite + scalar, and input_min<=input_max, output_min<=output_max.
+  Raises ValueError otherwise. No-op if ``cb`` is None."""
   if cb is None:
     return
+  vals = {}
   for nm in ("input_min", "input_max", "output_min", "output_max"):
     v = getattr(cb, nm).value
     if getattr(v, "shape", ()) not in ((), (1,)):
@@ -136,6 +138,57 @@ def validate_clip_bounds(cb, where=""):
       raise ValueError(f"Gemma4 vision clip bound '{nm}'{(' @ '+where) if where else ''} = {fv} is non-finite "
                        f"(missing/NaN/Inf). use_clipped_linears_for_vit=True declares a FINITE clipped model; "
                        f"refusing to fall back to an identity clamp.")
+    vals[nm] = fv
+  if vals["input_min"] > vals["input_max"]:
+    raise ValueError(f"Gemma4 vision clip bound{(' @ '+where) if where else ''}: input_min={vals['input_min']} "
+                     f"> input_max={vals['input_max']} (a clamp with min>max would empty the interval).")
+  if vals["output_min"] > vals["output_max"]:
+    raise ValueError(f"Gemma4 vision clip bound{(' @ '+where) if where else ''}: output_min={vals['output_min']} "
+                     f"> output_max={vals['output_max']} (a clamp with min>max would empty the interval).")
+
+
+# Expected clipped-module accounting for a Gemma-4 vision tower: 16 encoder blocks x 7 projections
+# (attention q/k/v/o = 4 modules holding q/k/v/o clip states + mlp gate/up/down = 3) -> but the clip
+# STATE lives on 2 modules per block (the Gemma4Attention holding q/k/v/o_clip, and the
+# Gemma4ClippedMlpBlock holding gate/up/down_clip). The scalar bound count is the invariant that matters:
+# 16 blocks x 7 projections x 4 bounds = 448. We validate the projection-level clip states (16x7 = 112).
+EXPECTED_CLIP_PROJECTIONS = 112   # 16 blocks x 7 projections (q,k,v,o,gate,up,down)
+EXPECTED_CLIP_BOUNDS = 448        # 112 projections x 4 scalar bounds
+
+
+def validate_all_vision_clip_bounds(model, *, expected_projections=EXPECTED_CLIP_PROJECTIONS,
+                                    expected_bounds=EXPECTED_CLIP_BOUNDS):
+  """Post-checkpoint-load, pre-first-JIT validation of ALL Gemma-4 vision clip bounds. Fail-closed.
+
+  Walks the model graph, collects every ``_ClipBounds`` module (each carries the 4 scalars for one
+  projection), validates each (finite, scalar, min<=max), and asserts EXACT counts:
+  ``expected_projections`` clip-state modules and ``expected_bounds`` scalar leaves. Raises ValueError on
+  any deficiency. The exact-count check is what prevents a traversal that accidentally finds zero modules
+  from silently "passing".
+  """
+  from flax import nnx  # pylint: disable=import-outside-toplevel
+  n_proj = 0
+  n_bounds = 0
+  for path, mod in nnx.iter_graph(model):
+    # A clip-state module has exactly the four bound leaves.
+    if (hasattr(mod, "input_min") and hasattr(mod, "input_max")
+        and hasattr(mod, "output_min") and hasattr(mod, "output_max")
+        and not isinstance(mod, (int, float))):
+      where = "/".join(str(getattr(p, "key", p)) for p in path) if isinstance(path, (list, tuple)) else str(path)
+      validate_clip_bounds(mod, where)
+      n_proj += 1
+      n_bounds += 4
+  if n_proj != expected_projections:
+    raise ValueError(
+        f"Gemma-4 vision clip validation found {n_proj} clip-state modules, expected exactly "
+        f"{expected_projections} (16 blocks x 7 projections). A mismatch means the clip bounds were not "
+        f"loaded/mapped correctly; refusing to run with a partially-clipped vision tower."
+    )
+  if n_bounds != expected_bounds:
+    raise ValueError(
+        f"Gemma-4 vision clip validation found {n_bounds} clip-bound scalars, expected exactly {expected_bounds}."
+    )
+  return n_proj, n_bounds
 
 
 
@@ -545,6 +598,14 @@ def enable_vision_clip_bounds(self):
     """
     if getattr(self, "_use_clipped_linears", False):
       return
+    # Fail-closed: the clipped path clamps q/k/v independently, which requires SEPARATE q/k/v projections.
+    # A fused QKV projection would apply a single clamp and silently bypass the per-projection clip semantics.
+    if bool(getattr(self.config, "fused_qkv", False)):
+      raise ValueError(
+          "Gemma-4 vision clipped-linears require fused_qkv=False: the checkpoint carries distinct "
+          "q/k/v activation clip bounds that must be applied per-projection. Refusing to silently clip a "
+          "fused QKV projection."
+      )
     self._use_clipped_linears = True
     self.q_clip = _make_clip_state()
     self.k_clip = _make_clip_state()
diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py
index a2a1403125..9e4bb54dce 100644
--- a/src/maxtext/utils/model_creation_utils.py
+++ b/src/maxtext/utils/model_creation_utils.py
@@ -1211,6 +1211,14 @@ def _walk_align(ckpt, model_arr, axes):
             "Please ensure the checkpoint format matches the scan_layers setting."
         )
 
+    # Gemma-4 vision clipped-linears: fail-closed post-load validation BEFORE the first JIT forward/train.
+    # Asserts exactly 112 clip-state modules / 448 finite, correctly-ordered scalar bounds. A missing/NaN/
+    # unordered bound or wrong module count raises here rather than silently degrading to an identity clamp.
+    if bool(getattr(config, "use_clipped_linears_for_vit", False)) and bool(getattr(config, "use_multimodal", False)):
+      from maxtext.models import gemma4_vision  # pylint: disable=import-outside-toplevel
+      n_proj, n_bounds = gemma4_vision.validate_all_vision_clip_bounds(model)
+      max_logging.log(f"Gemma-4 vision clip validation OK: {n_proj} clip-state modules, {n_bounds} scalar bounds.")
+
     if wrap_with_tunix_adapter:
       with mesh:
         use_no_op_mappings = "maxtext_config" in config.vllm_additional_config

From d9ac7522872fcecc90ab08f1bd84c594debb784d Mon Sep 17 00:00:00 2001
From: dengcchi <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 08:07:46 -0700
Subject: [PATCH 07/13] M1-M5 in-tree tests:
 tests/unit/gemma4_clip_linears_test.py (13 tests, all pass)

CI-grade unit tests covering the merge-critical invariants: clip math bit-exact vs jnp.clip,
dtype-preserving, exact no-op when disabled, stop_gradient => zero grad wrt bound, NaN/Inf/min>max
hard-fail, freeze-mask polarity+count, expected 112/448 count constants, public symbols present.
Runs on CPU in grr-maxtext:latest (13 passed).
---
 tests/unit/gemma4_clip_linears_test.py | 145 +++++++++++++++++++++++++
 1 file changed, 145 insertions(+)
 create mode 100644 tests/unit/gemma4_clip_linears_test.py

diff --git a/tests/unit/gemma4_clip_linears_test.py b/tests/unit/gemma4_clip_linears_test.py
new file mode 100644
index 0000000000..a8174a228b
--- /dev/null
+++ b/tests/unit/gemma4_clip_linears_test.py
@@ -0,0 +1,145 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""In-tree tests for the Gemma-4 vision clipped-linears (PR #4790).
+
+Pure mechanism + validation tests: no checkpoint, no TPU. Cover the merge-critical invariants:
+clip math vs jnp.clip, dtype/no-op, NaN/Inf/ordering hard-fail, exact 112/448 counts, optimizer
+freeze mask polarity, and fail-closed contract guards.
+"""
+
+import unittest
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from maxtext.models import gemma4_vision as gv
+
+jax.config.update("jax_platform_name", "cpu")
+
+
+def _finite_cb(in_lo=-2.5, in_hi=2.5, out_lo=-4.0, out_hi=4.0):
+  cb = gv._make_clip_state()
+  cb.input_min.value = jnp.asarray(in_lo, jnp.float32)
+  cb.input_max.value = jnp.asarray(in_hi, jnp.float32)
+  cb.output_min.value = jnp.asarray(out_lo, jnp.float32)
+  cb.output_max.value = jnp.asarray(out_hi, jnp.float32)
+  return cb
+
+
+class ClipMathTest(unittest.TestCase):
+
+  def test_clip_in_matches_jnp_clip(self):
+    cb = _finite_cb()
+    x = (jax.random.normal(jax.random.PRNGKey(0), (5, 17)) * 6.0).astype(jnp.float32)
+    self.assertTrue(bool(jnp.array_equal(gv._clip_in(x, cb), jnp.clip(x, -2.5, 2.5))))
+
+  def test_clip_out_matches_jnp_clip(self):
+    cb = _finite_cb()
+    x = (jax.random.normal(jax.random.PRNGKey(1), (5, 17)) * 6.0).astype(jnp.float32)
+    self.assertTrue(bool(jnp.array_equal(gv._clip_out(x, cb), jnp.clip(x, -4.0, 4.0))))
+
+  def test_none_is_exact_noop(self):
+    x = jax.random.normal(jax.random.PRNGKey(2), (3, 8)).astype(jnp.float32)
+    self.assertTrue(bool(jnp.array_equal(gv._clip_in(x, None), x)))
+    self.assertTrue(bool(jnp.array_equal(gv._clip_out(x, None), x)))
+
+  def test_dtype_preserved(self):
+    cb = _finite_cb()
+    x = jax.random.normal(jax.random.PRNGKey(3), (4, 4)).astype(jnp.bfloat16)
+    self.assertEqual(gv._clip_in(x, cb).dtype, jnp.bfloat16)
+
+  def test_bounds_are_stop_gradient(self):
+    # The clamp bound is read through jax.lax.stop_gradient at the use-site, so gradient wrt the bound
+    # value is exactly zero even when the clamp is active. Probe _clip_bound_value directly with a
+    # lightweight bound holder to avoid mutating an nnx Param inside a traced function.
+    class _B:
+      def __init__(self, v):
+        self.value = v
+
+    def loss(bound_val):
+      x = jnp.ones((4,), jnp.float32) * 10.0  # above the bound -> clamp active
+      hi = gv._clip_bound_value(_B(bound_val), jnp.float32)
+      return jnp.sum(jnp.clip(x, -100.0, hi))
+
+    g = jax.grad(loss)(jnp.asarray(2.5, jnp.float32))
+    self.assertEqual(float(g), 0.0)
+
+
+class ValidateBoundsTest(unittest.TestCase):
+
+  def test_accepts_finite_ordered(self):
+    gv.validate_clip_bounds(_finite_cb(), "ok")  # no raise
+
+  def test_rejects_nan_sentinel(self):
+    with self.assertRaises(ValueError):
+      gv.validate_clip_bounds(gv._make_clip_state(), "nan")
+
+  def test_rejects_inf(self):
+    cb = _finite_cb()
+    cb.input_max.value = jnp.asarray(np.inf, jnp.float32)
+    with self.assertRaises(ValueError):
+      gv.validate_clip_bounds(cb, "inf")
+
+  def test_rejects_min_gt_max(self):
+    cb = _finite_cb(in_lo=3.0, in_hi=-3.0)  # input_min > input_max
+    with self.assertRaises(ValueError):
+      gv.validate_clip_bounds(cb, "order")
+
+  def test_none_is_noop(self):
+    self.assertIsNone(gv.validate_clip_bounds(None))
+
+
+class FreezeMaskTest(unittest.TestCase):
+
+  def test_freeze_mask_polarity_and_count(self):
+    tree = {
+        "vision_encoder": {
+            "layer_0": {
+                "attention": {
+                    "q_clip": {"input_min": 1.0, "input_max": 1.0, "output_min": 1.0, "output_max": 1.0},
+                    "query": {"kernel": 1.0},
+                },
+                "mlp": {"gate_clip": {"input_min": 1.0, "output_max": 1.0}, "wi_0": {"kernel": 1.0}},
+            }
+        },
+        "decoder": {"layer_0": {"norm": {"scale": 1.0}}},
+    }
+    mask = gv.clip_optimizer_freeze_mask(tree)
+    flat = jax.tree_util.tree_flatten_with_path(mask)[0]
+    n_clip = 0
+    for path, v in flat:
+      is_clip = gv._is_clip_bound_path(path)
+      # clip leaves -> False (frozen); everything else -> True (trainable)
+      self.assertEqual(bool(v), (not is_clip))
+      if is_clip:
+        n_clip += 1
+    self.assertEqual(n_clip, 6)  # q_clip(4) + gate_clip(2)
+
+
+class SymbolsTest(unittest.TestCase):
+
+  def test_expected_counts_constants(self):
+    self.assertEqual(gv.EXPECTED_CLIP_PROJECTIONS, 112)
+    self.assertEqual(gv.EXPECTED_CLIP_BOUNDS, 448)
+
+  def test_public_symbols_present(self):
+    for sym in ("validate_all_vision_clip_bounds", "clip_optimizer_freeze_mask",
+                "Gemma4ClippedMlpBlock", "Gemma4Attention"):
+      self.assertTrue(hasattr(gv, sym), sym)
+
+
+if __name__ == "__main__":
+  unittest.main()

From 5bbca57227226ddd88d1eb4e7f8d083456516a40 Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 08:30:38 -0700
Subject: [PATCH 08/13] M1/M2/M6: stock-trainer position threading + PLE enum +
 MM contract gate

M1: thread image_position_ids through the stock loss_fn (train.py) for both
    Linen and NNX paths, guarded on presence in the batch (the stock fixed-grid
    Gemma-4 processor does not emit it -> legacy all-valid path, correct for its
    contract; native-resolution/pan-and-scan processors do -> Option-S path).
M2: make ple_pad_mode a validated enum (hard-fail unknown); 'identity' documented
    as the HF-faithful default (matches Transformers 5.9.0 get_per_layer_inputs,
    which ignores inputs_embeds when input_ids is given); wire 'both' as a real
    (HF-divergent, ablation-only) mode by threading shared_embedding into the
    Gemma-4-small PLE path -- removing the previously-dead hasattr(self,...) branch.
M6: Gemma-4 E2B/E4B multimodal static contract gate in config validation:
    enum PLE mode, non-negative token ids (single source of truth), causal image
    spans (use_bidirectional_image_attn=False), fused_qkv/fused_mlp=False.
    Pure-Linen Gemma-4-small MM path now fail-closed (PLE substitution is NNX-only).

No behavior change for non-Gemma-4 models or text-only Gemma-4.
---
 src/maxtext/configs/types.py            | 59 ++++++++++++++++++++++++-
 src/maxtext/layers/decoders.py          | 12 +++++
 src/maxtext/layers/nnx_decoders.py      | 17 ++++++-
 src/maxtext/trainers/pre_train/train.py | 13 ++++++
 4 files changed, 97 insertions(+), 4 deletions(-)

diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 712c042c85..31aee4192f 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2164,8 +2164,12 @@ class MultimodalGeneral(BaseModel):
   ple_pad_mode: str = Field(
       "identity",
       description=(
-          "PLE pad-substitution scope when ple_pad_substitute_image_rows=True: 'identity' (token-id path only) "
-          "or 'both' (also substitute the pad embedding in the context path)."
+          "PLE pad-substitution scope when ple_pad_substitute_image_rows=True. 'identity' (DEFAULT, "
+          "HF-faithful): map image placeholder rows -> pad in the token-identity PLE path only; the context/"
+          "projection path keeps the merged image features (matches HF Transformers 5.9.0 "
+          "Gemma4ForConditionalGeneration, where get_per_layer_inputs ignores inputs_embeds when input_ids is "
+          "provided). 'both': additionally overwrite the context path with the pad embedding — this is "
+          "HF-DIVERGENT and provided only for ablation. Validated as an enum; unknown values hard-fail."
       ),
   )
   image_placeholder_token_id: int = Field(
@@ -3617,6 +3621,57 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
             "(the vision encoder ships per-projection activation clip bounds; without them the "
             "image span diverges). Set use_clipped_linears_for_vit=True to enable image inputs."
         )
+      # ---- Gemma-4 E2B/E4B multimodal STATIC contract gate (fail-closed) ----
+      # These invariants encode the semantics validated against pinned HF Transformers 5.9.0
+      # (Gemma4ForConditionalGeneration). Any deviation silently corrupts image/post-image logits, so we
+      # refuse to build the model rather than degrade to a wrong-but-runnable path.
+      if self.model_name in ("gemma4-e2b", "gemma4-e4b"):
+        # (a) PLE pad-substitution mode must be a known value. HF maps image placeholder tokens -> pad in the
+        #     token-identity PLE path only (the context path keeps the merged image features); that is
+        #     "identity". "both" additionally overwrites the context path, which is HF-DIVERGENT. We keep the
+        #     knob for experimentation but hard-fail unknown/typo values instead of silently defaulting.
+        _valid_ple_modes = ("identity", "both")
+        if str(self.ple_pad_mode) not in _valid_ple_modes:
+          raise ValueError(
+              f"ple_pad_mode='{self.ple_pad_mode}' is not one of {_valid_ple_modes}. "
+              f"For Gemma-4 E2B/E4B the HF-faithful contract is 'identity' (token-identity PLE path maps "
+              f"image rows -> pad; context path keeps merged image features). 'both' is HF-divergent and "
+              f"provided only for ablation. Refusing to run with an unrecognized PLE mode."
+          )
+        # (b) The image placeholder id and PLE pad id are semantic constants tied to the tokenizer/model. We
+        #     require them to be set explicitly (single source of truth: the model yml derived from HF config)
+        #     so a silent hidden default cannot mask a tokenizer mismatch.
+        if self.ple_pad_substitute_image_rows:
+          if int(self.image_placeholder_token_id) < 0:
+            raise ValueError(
+                "image_placeholder_token_id must be a valid non-negative token id when "
+                "ple_pad_substitute_image_rows=True (derive it from the model/tokenizer config)."
+            )
+          if int(self.ple_pad_token_id) < 0:
+            raise ValueError(
+                "ple_pad_token_id must be a valid non-negative token id when "
+                "ple_pad_substitute_image_rows=True (Gemma-4 E2B text_config.pad_token_id=0)."
+            )
+        # (c) Gemma-4 E2B/E4B image spans are CAUSAL. Bidirectional image attention is a Gemma-3 / 26B / 31B
+        #     feature and would change the attention pattern for E2B/E4B.
+        if bool(getattr(self, "use_bidirectional_image_attn", False)):
+          raise ValueError(
+              f"{self.model_name} uses CAUSAL image spans; use_bidirectional_image_attn must be False. "
+              "Bidirectional image attention is for Gemma-3 / gemma4-26b / gemma4-31b."
+          )
+        # (d) The clipped path clamps q/k/v and gate/up/down per-projection; a fused QKV or fused MLP would
+        #     apply a single clamp and silently bypass the per-projection clip semantics. Fail closed here
+        #     (defense in depth alongside the runtime guards in gemma4_vision).
+        if bool(getattr(self, "fused_qkv", False)):
+          raise ValueError(
+              f"{self.model_name} multimodal clipped-linears require fused_qkv=False "
+              "(distinct q/k/v activation clip bounds must be applied per-projection)."
+          )
+        if bool(getattr(self, "fused_mlp", False)):
+          raise ValueError(
+              f"{self.model_name} multimodal clipped-linears require fused_mlp=False "
+              "(distinct gate/up/down activation clip bounds must be applied per-projection)."
+          )
       valid_mm_models = (
           "gemma3-4b",
           "gemma3-12b",
diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py
index 42753eb752..fad5eefa0b 100644
--- a/src/maxtext/layers/decoders.py
+++ b/src/maxtext/layers/decoders.py
@@ -1679,6 +1679,18 @@ def _apply_gemma4_small_layers(
 
     per_layer_inputs = None
     if cfg.hidden_size_per_layer_input > 0 and cfg.vocab_size_per_layer_input > 0:
+      # Fail-closed: the pure-Linen Gemma-4-small decoder does NOT implement the PLE image-row pad
+      # substitution (the HF-faithful mapping of image placeholder rows -> pad in the token-identity PLE
+      # path). That logic lives in the NNX decoder (nnx_decoders._apply_gemma4_small_layers) and is exercised
+      # by the supported enable_nnx=True path. Running E2B/E4B multimodal through this pure-Linen path would
+      # silently diverge on the image-span / post-image logits, so we refuse rather than degrade quietly.
+      if bool(getattr(cfg, "ple_pad_substitute_image_rows", False)) and multimodal_input is not None:
+        raise NotImplementedError(
+            "Gemma-4 E2B/E4B multimodal PLE pad-substitution is only implemented on the NNX decoder path "
+            "(enable_nnx=True, the default). The pure-Linen decoder does not perform the image-row PLE "
+            "substitution and would silently diverge from the HF reference. Set enable_nnx=True (default) to "
+            "run Gemma-4 E2B/E4B multimodal."
+        )
       per_layer_inputs = gemma4_small.PLEToLinen(
           config=cfg,
           mesh=mesh,
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 51b32ea112..68412bb6fc 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -1754,6 +1754,7 @@ def __call__(
             attention_metadata=attention_metadata,
             previous_chunk=previous_chunk,
             slot=slot,
+            shared_embedding=shared_embedding,
         )
       elif cfg.scan_layers:
         if self.is_deepseek:
@@ -2227,6 +2228,7 @@ def _apply_gemma4_small_layers(
       attention_metadata=None,
       previous_chunk=None,
       slot=None,
+      shared_embedding=None,
   ):
     """Apply Gemma 4 small (E2B/E4B) decoder layers (pure-NNX)."""
     cfg = self.config
@@ -2252,8 +2254,19 @@ def _apply_gemma4_small_layers(
         _pad_id = int(getattr(cfg, "ple_pad_token_id", 0))
         _img_row = decoder_input_tokens.astype(jnp.int32) == _img_id
         ple_tokens = jnp.where(_img_row, _pad_id, decoder_input_tokens.astype(jnp.int32))
-        if str(getattr(cfg, "ple_pad_mode", "identity")) == "both" and hasattr(self, "shared_embedding"):
-          _pad_vec = self.shared_embedding(jnp.full_like(decoder_input_tokens, _pad_id).astype(jnp.int32))
+        # ple_pad_mode: 'identity' (HF-faithful default) substitutes ONLY the token-identity path (above),
+        # leaving the context/projection path reading the merged image features — this matches HF
+        # Transformers 5.9.0 (get_per_layer_inputs ignores inputs_embeds when input_ids is given). The
+        # 'both' mode additionally overwrites the context path with the pad embedding; this is HF-DIVERGENT
+        # and provided only for ablation (see M9 ablation matrix). Enum-validated in configs/types.py.
+        if str(getattr(cfg, "ple_pad_mode", "identity")) == "both":
+          if shared_embedding is None:
+            raise ValueError(
+                "ple_pad_mode='both' requires the shared token embedding to compute the pad context "
+                "vector, but shared_embedding was not threaded into the Gemma-4-small PLE path. This is "
+                "an HF-divergent ablation mode; use the default 'identity' for HF-faithful behavior."
+            )
+          _pad_vec = shared_embedding(jnp.full_like(decoder_input_tokens, _pad_id).astype(jnp.int32), model_mode=model_mode)
           ple_context = jnp.where(_img_row[..., None], _pad_vec, y)
       per_layer_inputs = self.per_layer_embedder(ple_tokens, ple_context)
 
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 0906c3e786..d191426072 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -148,6 +148,14 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
         decoder_segment_ids=data["inputs_segmentation"],
         encoder_images=data["images"] if config.use_multimodal else None,
         encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None,
+        # Gemma-4 padded-patch (Option-S) contract: when the data pipeline produces per-patch image
+        # position ids (native-resolution / pan-and-scan processors), thread them to the vision encoder
+        # so pad patches are masked and pooled by real coordinates. The stock fixed-grid Gemma-4 processor
+        # produces all-valid patches and does NOT emit this key, so the legacy all-valid path is used
+        # (correct for its contract). See gemma4_vision.Gemma4VisionEncoderLayer for the fail-closed guard.
+        encoder_image_position_ids=(
+            data["image_position_ids"] if config.use_multimodal and "image_position_ids" in data else None
+        ),
         enable_dropout=config.enable_dropout if is_train else False,
         rngs={"dropout": rng1, "params": aqt_rng},  # pyrefly: ignore[bad-argument-type]
         mutable=mutable_collections,
@@ -197,6 +205,11 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr
         decoder_segment_ids=data["inputs_segmentation"],
         encoder_images=data["images"] if config.use_multimodal else None,
         encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None,
+        # See the Linen branch above: thread per-patch image position ids when the data pipeline emits
+        # them (Gemma-4 Option-S padded-patch contract); None for the stock fixed-grid all-valid processor.
+        encoder_image_position_ids=(
+            data["image_position_ids"] if config.use_multimodal and "image_position_ids" in data else None
+        ),
         enable_dropout=config.enable_dropout if is_train else False,
         decoder_target_tokens=data["targets"],
         decoder_target_mask=data["targets_segmentation"],

From c45b22426878b080d1b4742f36b104cb3eed18ae Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 08:57:18 -0700
Subject: [PATCH 09/13] M2/M5/M6 in-tree test: Gemma-4 MM static contract gate
 (enum/causal/fused/clipped)

---
 tests/unit/gemma4_mm_contract_gate_test.py | 83 ++++++++++++++++++++++
 1 file changed, 83 insertions(+)
 create mode 100644 tests/unit/gemma4_mm_contract_gate_test.py

diff --git a/tests/unit/gemma4_mm_contract_gate_test.py b/tests/unit/gemma4_mm_contract_gate_test.py
new file mode 100644
index 0000000000..06b49e5bd0
--- /dev/null
+++ b/tests/unit/gemma4_mm_contract_gate_test.py
@@ -0,0 +1,83 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""M2/M5/M6: Gemma-4 E2B/E4B multimodal STATIC contract gate (configs/types.py).
+
+Asserts the config validator ACCEPTS a valid E2B multimodal config and HARD-FAILS on each
+contract violation: unknown PLE mode (M2 enum), bidirectional image attention (causal M6),
+fused_qkv / fused_mlp (M5), and clipped-linears disabled. These encode the semantics verified
+against pinned HF Transformers 5.9.0 and must fail closed rather than silently degrade.
+"""
+import os
+import unittest
+
+from maxtext.configs import pyconfig
+from maxtext.utils.globals import MAXTEXT_PKG_DIR
+
+
+def _init(overrides):
+  base_yml = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml")
+  argv = [
+      "test",
+      base_yml,
+      "model_name=gemma4-e2b",
+      "use_multimodal=true",
+      "override_model_config=true",
+      "tokenizer_path=/tmp/unused",
+      "per_device_batch_size=1",
+      "scan_layers=false",
+      "enable_checkpointing=false",
+      "run_name=gate_test",
+      "steps=1",
+      "skip_jax_distributed_system=true",
+  ] + overrides
+  return pyconfig.initialize(argv)
+
+
+class Gemma4MmContractGateTest(unittest.TestCase):
+
+  def test_valid_e2b_mm_config_builds(self):
+    # Clipped linears on, identity PLE, causal image spans, fused off -> the HF-faithful contract.
+    cfg = _init(["use_clipped_linears_for_vit=true"])
+    self.assertEqual(cfg.model_name, "gemma4-e2b")
+    self.assertTrue(cfg.use_multimodal)
+    self.assertEqual(str(cfg.ple_pad_mode), "identity")
+
+  def test_unknown_ple_pad_mode_hard_fails(self):
+    with self.assertRaises(Exception):
+      _init(["use_clipped_linears_for_vit=true", "ple_pad_mode=bogus"])
+
+  def test_both_ple_pad_mode_is_valid_ablation_enum(self):
+    cfg = _init(["use_clipped_linears_for_vit=true", "ple_pad_mode=both"])
+    self.assertEqual(str(cfg.ple_pad_mode), "both")
+
+  def test_bidirectional_image_attn_hard_fails(self):
+    # Gemma-4 E2B/E4B image spans are causal; bidirectional is a Gemma-3 / 26B / 31B feature.
+    with self.assertRaises(Exception):
+      _init(["use_clipped_linears_for_vit=true", "use_bidirectional_image_attn=true"])
+
+  def test_fused_qkv_hard_fails(self):
+    with self.assertRaises(Exception):
+      _init(["use_clipped_linears_for_vit=true", "fused_qkv=true"])
+
+  def test_fused_mlp_hard_fails(self):
+    with self.assertRaises(Exception):
+      _init(["use_clipped_linears_for_vit=true", "fused_mlp=true"])
+
+  def test_clipped_linears_off_under_mm_hard_fails(self):
+    with self.assertRaises(Exception):
+      _init(["use_clipped_linears_for_vit=false"])
+
+
+if __name__ == "__main__":
+  unittest.main()

From dba5cf773b35f6c264cba13e58c7c38e394f9bda Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 09:00:28 -0700
Subject: [PATCH 10/13] M7: Option-S sentinel integrity + placeholder==pooled
 validators + 5 unit tests

---
 src/maxtext/models/gemma4_vision.py    | 62 ++++++++++++++++++++++++++
 tests/unit/gemma4_clip_linears_test.py | 37 +++++++++++++++
 2 files changed, 99 insertions(+)

diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index 541403296e..4456ec5d1d 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -191,8 +191,70 @@ def validate_all_vision_clip_bounds(model, *, expected_projections=EXPECTED_CLIP
   return n_proj, n_bounds
 
 
+# Sentinel value marking a padded (phantom) patch position in the Option-S padded-patch contract.
+POSITIONS_PAD_VALUE = -1
 
 
+def option_s_position_sentinel_ok(positions_xy):
+  """Return (all_rows_valid_or_full_sentinel, n_bad_rows) for an Option-S position tensor.
+
+  The padded-patch contract requires every per-patch position row to be EXACTLY one of:
+    * [x, y] with both coordinates >= 0 (a valid patch), or
+    * [POSITIONS_PAD_VALUE, POSITIONS_PAD_VALUE] (a fully-padded sentinel).
+  A MIXED row such as [-1, y] or [x, -1] is malformed: it would be classified as "valid" by the
+  ``(pos == -1).all(axis=-1)`` pad test (silently attended-to as a real patch) yet carry a negative
+  coordinate into pooling. This helper flags such rows so callers can fail closed rather than
+  silently mis-handle them. Pure (numpy/jax-eager friendly); does not allocate on-device state.
+
+  Args:
+    positions_xy: int array [..., 2] of per-patch (x, y) positions.
+  Returns:
+    (ok: bool, n_bad: int) where ok is True iff no malformed (mixed-sentinel) row exists.
+  """
+  import numpy as _np  # local import: this is called eagerly (host) for validation, not in the JIT forward
+  p = _np.asarray(positions_xy)
+  if p.shape[-1] != 2:
+    raise ValueError(f"option_s positions must have last dim 2 (x,y); got shape {p.shape}")
+  is_neg = p == POSITIONS_PAD_VALUE
+  any_neg = is_neg.any(axis=-1)
+  all_neg = is_neg.all(axis=-1)
+  # A row is malformed iff it has SOME sentinel coord but is not FULLY sentinel.
+  bad = _np.logical_and(any_neg, _np.logical_not(all_neg))
+  # Also reject any negative coordinate other than the exact sentinel (e.g. -2), which is never valid.
+  other_neg = _np.logical_and(p < 0, p != POSITIONS_PAD_VALUE).any(axis=-1)
+  bad = _np.logical_or(bad, other_neg)
+  n_bad = int(bad.sum())
+  return (n_bad == 0), n_bad
+
+
+def validate_option_s_positions(positions_xy, where=""):
+  """Hard-fail if any Option-S position row is a malformed mixed sentinel. No-op if positions is None."""
+  if positions_xy is None:
+    return
+  ok, n_bad = option_s_position_sentinel_ok(positions_xy)
+  if not ok:
+    raise ValueError(
+        f"Gemma-4 Option-S position sentinel integrity{(' @ '+where) if where else ''}: found {n_bad} malformed "
+        f"row(s). Every patch position must be [x,y] with both coords >= 0, or [-1,-1] (full sentinel). Mixed "
+        f"rows like [-1, y] / [x, -1] (or other negatives) would be silently treated as valid patches and "
+        f"corrupt pooling; refusing to run."
+    )
+
+
+def option_s_pooled_matches_placeholder(image_mask, n_placeholder):
+  """Return (ok, n_valid_pooled) checking the count of valid pooled image tokens == n_placeholder.
+
+  ``image_mask`` is the [..., K] validity mask returned by the vision pooler (True where a pooled slot
+  received at least one real patch). ``n_placeholder`` is the number of image placeholder positions in the
+  decoder sequence for the SAME sample. The Option-S contract requires these to be equal so that the pooled
+  image tokens scatter 1:1 onto the placeholder rows; a mismatch means the merge would drop or duplicate
+  image tokens regardless of padded-embedding ordering.
+  """
+  import numpy as _np
+  m = _np.asarray(image_mask)
+  n_valid = int(_np.asarray(m).astype(bool).sum())
+  return (n_valid == int(n_placeholder)), n_valid
+
 
 def factorized_posemb(posemb: jax.Array, positions_xy: jax.Array, precision) -> jax.Array:
   """Computes factorized position embedding from (x, y) coordinates.
diff --git a/tests/unit/gemma4_clip_linears_test.py b/tests/unit/gemma4_clip_linears_test.py
index a8174a228b..fd2c08f766 100644
--- a/tests/unit/gemma4_clip_linears_test.py
+++ b/tests/unit/gemma4_clip_linears_test.py
@@ -143,3 +143,40 @@ def test_public_symbols_present(self):
 
 if __name__ == "__main__":
   unittest.main()
+
+
+class OptionSInvariantsTest(unittest.TestCase):
+  """M7: Option-S sentinel integrity + placeholder==pooled count invariants (pure validators)."""
+
+  def test_valid_positions_pass(self):
+    pos = np.array([[[0, 0], [1, 0], [-1, -1], [-1, -1]]], dtype=np.int32)  # 2 valid + 2 full-sentinel
+    ok, n_bad = gv.option_s_position_sentinel_ok(pos)
+    self.assertTrue(ok)
+    self.assertEqual(n_bad, 0)
+    gv.validate_option_s_positions(pos)  # must not raise
+
+  def test_mixed_sentinel_row_fails(self):
+    for bad_row in ([-1, 5], [5, -1]):
+      pos = np.array([[[0, 0], bad_row, [-1, -1]]], dtype=np.int32)
+      ok, n_bad = gv.option_s_position_sentinel_ok(pos)
+      self.assertFalse(ok, f"mixed row {bad_row} should be flagged")
+      self.assertEqual(n_bad, 1)
+      with self.assertRaises(ValueError):
+        gv.validate_option_s_positions(pos, where="test")
+
+  def test_other_negative_coord_fails(self):
+    pos = np.array([[[0, 0], [-2, -2], [-1, -1]]], dtype=np.int32)  # -2 is not the sentinel
+    ok, n_bad = gv.option_s_position_sentinel_ok(pos)
+    self.assertFalse(ok)
+
+  def test_pooled_matches_placeholder(self):
+    mask = np.array([True, True, True, False, False])  # 3 valid pooled tokens
+    ok, n_valid = gv.option_s_pooled_matches_placeholder(mask, 3)
+    self.assertTrue(ok)
+    self.assertEqual(n_valid, 3)
+
+  def test_pooled_placeholder_mismatch_detected(self):
+    mask = np.array([True, True, True, False, False])  # 3 valid
+    ok, n_valid = gv.option_s_pooled_matches_placeholder(mask, 4)  # placeholder count 4 != 3
+    self.assertFalse(ok)
+    self.assertEqual(n_valid, 3)

From 8fef20ba65d3ec42691ae1a302ae191a685cfd57 Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 09:07:25 -0700
Subject: [PATCH 11/13] M7.4: sanitize padded-patch rows to zero before
 attention (isolate non-finite pad from valid outputs)

---
 src/maxtext/models/gemma4_vision.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/maxtext/models/gemma4_vision.py b/src/maxtext/models/gemma4_vision.py
index 4456ec5d1d..b65836b39c 100644
--- a/src/maxtext/models/gemma4_vision.py
+++ b/src/maxtext/models/gemma4_vision.py
@@ -1037,6 +1037,16 @@ def __call__(
     is_pad = (positions_xy == -1).all(axis=-1)  # [B*N, L]
     decoder_segment_ids = jnp.where(is_pad, 2, 1).astype(jnp.int32)
 
+    # Sanitize padded-patch rows to zero BEFORE attention. Segment masking already prevents pad patches
+    # from influencing valid patches' attention outputs, and pooling excludes pad positions — so for FINITE
+    # pad content the valid outputs are provably invariant (verified byte-identical under 1e6 garbage pad).
+    # However, attention computes q/k/v on ALL rows before masking, so a NON-FINITE (NaN/Inf) pad patch would
+    # produce NaN q/k/v and leak through the softmax normalization into valid rows (NaN*0=NaN). Zeroing pad
+    # rows here makes the isolation hold for arbitrary (incl. non-finite) pad content without changing any
+    # valid output. This is defense-in-depth for malformed processor output; the frozen processor emits finite
+    # pad, which was already invariant.
+    x = jnp.where(is_pad[..., None], jnp.zeros_like(x), x)
+
     for i in range(self.config.num_hidden_layers_for_vit):
       layer = getattr(self, f"layer_{i}")
       x = layer(

From 974ceeebcfebfc13456835dd3adffa2c7f02c9f8 Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 09:16:37 -0700
Subject: [PATCH 12/13] M8: fail-close packing for E2B/E4B multimodal (any
 mode) + packing-guard test

Packed multimodal is unsupported by the stock data pipeline; packing image spans risks cross-doc
image attention + PLE image-row substitution across segment boundaries. MaxText forbade this for
multimodal SFT only; extend to every training mode. + in-tree packing-guard test.
---
 src/maxtext/configs/types.py               | 11 +++++++++++
 tests/unit/gemma4_mm_contract_gate_test.py |  5 +++++
 2 files changed, 16 insertions(+)

diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 31aee4192f..69a3d29270 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -3672,6 +3672,17 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
               f"{self.model_name} multimodal clipped-linears require fused_mlp=False "
               "(distinct gate/up/down activation clip bounds must be applied per-projection)."
           )
+        # (e) Sequence packing is not supported for Gemma-4 E2B/E4B multimodal in ANY training mode. The stock
+        #     data pipeline does not pack image spans, and packing image-bearing examples together would risk
+        #     cross-document image attention and PLE image-row substitution bleeding across segment boundaries.
+        #     MaxText already forbids packing for multimodal SFT; extend that to every mode for these models
+        #     (fail closed rather than silently produce cross-doc image attention).
+        if bool(getattr(self, "packing", False)):
+          raise ValueError(
+              f"{self.model_name} multimodal does not support sequence packing (packing=True). The stock data "
+              "pipeline does not pack image spans; packing image-bearing examples risks cross-document image "
+              "attention and PLE image-row substitution across segment boundaries. Set packing=False."
+          )
       valid_mm_models = (
           "gemma3-4b",
           "gemma3-12b",
diff --git a/tests/unit/gemma4_mm_contract_gate_test.py b/tests/unit/gemma4_mm_contract_gate_test.py
index 06b49e5bd0..a60529079b 100644
--- a/tests/unit/gemma4_mm_contract_gate_test.py
+++ b/tests/unit/gemma4_mm_contract_gate_test.py
@@ -78,6 +78,11 @@ def test_clipped_linears_off_under_mm_hard_fails(self):
     with self.assertRaises(Exception):
       _init(["use_clipped_linears_for_vit=false"])
 
+  def test_packing_hard_fails(self):
+    # Packed multimodal is unsupported for E2B/E4B in any mode (cross-doc image attention risk).
+    with self.assertRaises(Exception):
+      _init(["use_clipped_linears_for_vit=true", "packing=true"])
+
 
 if __name__ == "__main__":
   unittest.main()

From 797842b1353df3fb134de67f47ad403385442114 Mon Sep 17 00:00:00 2001
From: Loki Chen <dengcchi@meta.com>
Date: Sun, 9 Aug 2026 09:21:27 -0700
Subject: [PATCH 13/13] E2B/E4B ymls: default packing=false (MM does not
 support packed image spans)

---
 src/maxtext/configs/models/gemma4-e2b.yml | 1 +
 src/maxtext/configs/models/gemma4-e4b.yml | 1 +
 2 files changed, 2 insertions(+)

diff --git a/src/maxtext/configs/models/gemma4-e2b.yml b/src/maxtext/configs/models/gemma4-e2b.yml
index c4a4b26c26..502e84737b 100644
--- a/src/maxtext/configs/models/gemma4-e2b.yml
+++ b/src/maxtext/configs/models/gemma4-e2b.yml
@@ -66,3 +66,4 @@ ple_pad_substitute_image_rows: True
 ple_pad_mode: "identity"
 image_placeholder_token_id: 258880
 ple_pad_token_id: 0
+packing: false  # Gemma-4 E2B/E4B multimodal does not support packed image spans (see configs/types.py gate)
diff --git a/src/maxtext/configs/models/gemma4-e4b.yml b/src/maxtext/configs/models/gemma4-e4b.yml
index cd752ae22b..d2cf5bdd7e 100644
--- a/src/maxtext/configs/models/gemma4-e4b.yml
+++ b/src/maxtext/configs/models/gemma4-e4b.yml
@@ -67,3 +67,4 @@ ple_pad_substitute_image_rows: True
 ple_pad_mode: "identity"
 image_placeholder_token_id: 258880
 ple_pad_token_id: 0
+packing: false  # Gemma-4 E2B/E4B multimodal does not support packed image spans (see configs/types.py gate)