From 6df741a32bead0f36d89be5a4cd088ce7acd8319 Mon Sep 17 00:00:00 2001
From: IMvision12 <gitesh.ch.0912@gmail.com>
Date: Fri, 28 Aug 2026 16:05:21 -0700
Subject: [PATCH 1/3] Fix TF SparseTensor issue accros the models

---
 tests/base/model_test_registry.py             | 33 +++++++++++++
 .../models/deepseek_v4/deepseek_v4_layers.py  | 16 +++++++
 zeromodels/models/gemma/gemma_layers.py       | 22 +++++++++
 zeromodels/models/gemma2/gemma2_layers.py     | 26 +++++++++++
 zeromodels/models/gemma3/gemma3_layers.py     | 24 ++++++++++
 zeromodels/models/gemma3n/gemma3n_layers.py   | 30 ++++++++++++
 .../models/glm4v/glm4v_vision_layers.py       | 46 +++++++++++++++++++
 zeromodels/models/glm5_moe/glm5_moe_layers.py | 33 +++++++++++++
 zeromodels/models/gpt_oss/gpt_oss_layers.py   | 12 +++++
 .../granite_speech5/granite_speech5_layers.py |  4 ++
 zeromodels/models/minimax/minimax_layers.py   | 24 ++++++++++
 .../minimax_m3_vl/minimax_m3_vl_layers.py     | 15 ++++++
 .../minimax_m3_vl/minimax_m3_vl_model.py      |  8 ++++
 zeromodels/models/qwen2/qwen2_layers.py       | 22 +++++++++
 .../models/qwen2_5_vl/qwen2_5_vl_layers.py    | 29 ++++++++++++
 .../models/qwen2_5_vl/qwen2_5_vl_model.py     |  9 ++++
 zeromodels/models/qwen2_vl/qwen2_vl_layers.py | 28 +++++++++++
 zeromodels/models/qwen2_vl/qwen2_vl_model.py  |  9 ++++
 zeromodels/models/qwen3_5/qwen3_5_layers.py   | 25 ++++++++++
 .../models/qwen3_5_moe/qwen3_5_moe_model.py   |  7 +++
 .../models/qwen3_next/qwen3_next_layers.py    | 29 ++++++++++++
 zeromodels/models/t5/t5_layers.py             | 36 +++++++++++++++
 22 files changed, 487 insertions(+)

diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py
index d3e132bb..e352be92 100644
--- a/tests/base/model_test_registry.py
+++ b/tests/base/model_test_registry.py
@@ -390,6 +390,39 @@
         "input_shape": (2, 32, 32, 3),
         "expected_output_shape": (2, 1000),
     },
+    "PvtImageClassify": {
+        "module": "zeromodels.models.pvt",
+        "model_cls": "PvtImageClassify",
+        "model_type": "classification",
+        "init_kwargs": {
+            "hidden_sizes": (8, 16, 40, 64),
+            "depths": (1, 1, 1, 1),
+            "num_attention_heads": (1, 2, 5, 8),
+            "sr_ratios": (8, 4, 2, 1),
+            "mlp_ratios": (2, 2, 2, 2),
+            "image_size": (32, 32, 3),
+            "num_classes": 1000,
+        },
+        "input_shape": (2, 32, 32, 3),
+        "expected_output_shape": (2, 1000),
+    },
+    "PvtV2ImageClassify": {
+        "module": "zeromodels.models.pvt_v2",
+        "model_cls": "PvtV2ImageClassify",
+        "model_type": "classification",
+        "init_kwargs": {
+            "hidden_sizes": (8, 16, 40, 64),
+            "depths": (1, 1, 1, 1),
+            "num_attention_heads": (1, 2, 5, 8),
+            "sr_ratios": (8, 4, 2, 1),
+            "mlp_ratios": (2, 2, 2, 2),
+            "linear_attention": False,
+            "image_size": (32, 32, 3),
+            "num_classes": 1000,
+        },
+        "input_shape": (2, 32, 32, 3),
+        "expected_output_shape": (2, 1000),
+    },
     "Res2NetImageClassify": {
         "module": "zeromodels.models.res2net",
         "model_cls": "Res2NetImageClassify",
diff --git a/zeromodels/models/deepseek_v4/deepseek_v4_layers.py b/zeromodels/models/deepseek_v4/deepseek_v4_layers.py
index ce2b0dee..8a67d58f 100644
--- a/zeromodels/models/deepseek_v4/deepseek_v4_layers.py
+++ b/zeromodels/models/deepseek_v4/deepseek_v4_layers.py
@@ -96,6 +96,12 @@ def __init__(self, embed_dim, mlp_dim, swiglu_limit=10.0, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(clamped_swiglu(self.gate(x), self.up(x), self.swiglu_limit))
 
@@ -942,6 +948,16 @@ def __init__(
             hc_mult, embed_dim, hc_sinkhorn_iters, hc_eps, name="ffn_hc"
         )
 
+    def build(self, input_shape):
+        collapsed_shape = tuple(input_shape[:-2]) + (self.embed_dim,)
+        self.attn_hc.build(input_shape)
+        self.ffn_hc.build(input_shape)
+        self.attention_norm.build(collapsed_shape)
+        self.attention.build(collapsed_shape)
+        self.mlp_norm.build(collapsed_shape)
+        self.mlp.build(collapsed_shape)
+        self.built = True
+
     def mix(self, streams, post, comb, sublayer_out):
         dtype = streams.dtype
         post = ops.cast(post, dtype)
diff --git a/zeromodels/models/gemma/gemma_layers.py b/zeromodels/models/gemma/gemma_layers.py
index a361ccc4..22904ac0 100644
--- a/zeromodels/models/gemma/gemma_layers.py
+++ b/zeromodels/models/gemma/gemma_layers.py
@@ -61,6 +61,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.gelu(self.gate(x), approximate=True) * self.up(x))
 
@@ -112,6 +118,13 @@ def __init__(self, embed_dim, num_heads, num_kv_heads, head_dim, **kwargs):
         self.value = layers.Dense(num_kv_heads * head_dim, use_bias=False, name="value")
         self.output_proj = layers.Dense(embed_dim, use_bias=False, name="output_proj")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.built = True
+
     def call(
         self,
         hidden_states,
@@ -250,6 +263,15 @@ def __init__(
         self.mlp_norm = GemmaRMSNorm(eps=norm_eps, name="mlp_norm")
         self.mlp = GemmaMLP(embed_dim, mlp_dim, name="mlp")
 
+    def build(self, input_shape):
+        # Explicit child builds so Keras never auto-builds via a call() trace, which
+        # runs GemmaRMSNorm.call() on a symbolic placeholder and fails on TF.
+        self.attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.mlp_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/gemma2/gemma2_layers.py b/zeromodels/models/gemma2/gemma2_layers.py
index 625edad5..af756c17 100644
--- a/zeromodels/models/gemma2/gemma2_layers.py
+++ b/zeromodels/models/gemma2/gemma2_layers.py
@@ -60,6 +60,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.gelu(self.gate(x), approximate=True) * self.up(x))
 
@@ -119,6 +125,13 @@ def __init__(
         self.value = layers.Dense(num_kv_heads * head_dim, use_bias=False, name="value")
         self.output_proj = layers.Dense(embed_dim, use_bias=False, name="output_proj")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.built = True
+
     def softcap(self, attn):
         if self.attn_logit_softcapping is None:
             return attn
@@ -290,6 +303,19 @@ def __init__(
             eps=norm_eps, name="post_feedforward_norm"
         )
 
+    def build(self, input_shape):
+        # Build children explicitly so Keras never auto-builds by tracing call(), which
+        # runs Gemma2RMSNorm.call() on a symbolic placeholder and fails on the TF
+        # backend (a SparseTensor from square/mean). Norms build here; attention/mlp are
+        # marked built and lazy-build their own children on the (proper) functional call.
+        self.attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.post_attention_norm.build(input_shape)
+        self.pre_feedforward_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.post_feedforward_norm.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/gemma3/gemma3_layers.py b/zeromodels/models/gemma3/gemma3_layers.py
index dbfffd64..5d40d1a9 100644
--- a/zeromodels/models/gemma3/gemma3_layers.py
+++ b/zeromodels/models/gemma3/gemma3_layers.py
@@ -62,6 +62,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.gelu(self.gate(x), approximate=True) * self.up(x))
 
@@ -124,6 +130,15 @@ def __init__(
         self.query_norm = Gemma3RMSNorm(eps=norm_eps, name="query_norm")
         self.key_norm = Gemma3RMSNorm(eps=norm_eps, name="key_norm")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.query_norm.build(tuple(input_shape[:-1]) + (self.num_heads, self.head_dim))
+        self.key_norm.build(tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim))
+        self.built = True
+
     def call(
         self,
         hidden_states,
@@ -290,6 +305,15 @@ def __init__(
             eps=norm_eps, name="post_feedforward_norm"
         )
 
+    def build(self, input_shape):
+        self.attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.post_attention_norm.build(input_shape)
+        self.pre_feedforward_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.post_feedforward_norm.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/gemma3n/gemma3n_layers.py b/zeromodels/models/gemma3n/gemma3n_layers.py
index f92074fb..27d396da 100644
--- a/zeromodels/models/gemma3n/gemma3n_layers.py
+++ b/zeromodels/models/gemma3n/gemma3n_layers.py
@@ -257,6 +257,18 @@ def __init__(
             num_heads * head_dim, use_bias=False, name="output_proj"
         )
 
+    def build(self, input_shape):
+        prefix = tuple(input_shape[:-1])
+        self.query.build(input_shape)
+        self.query_norm.build(prefix + (self.num_heads, self.head_dim))
+        if not self.is_kv_shared:
+            self.key.build(input_shape)
+            self.value.build(input_shape)
+            self.key_norm.build(prefix + (self.num_kv_heads, self.head_dim))
+            self.value_norm.build(prefix + (self.num_kv_heads, self.head_dim))
+        self.output_proj.build(prefix + (self.num_heads * self.head_dim,))
+        self.built = True
+
     def project_kv(self, hidden, cos, sin):
         b, s = int(hidden.shape[0]), int(hidden.shape[1])
         k = ops.reshape(self.key(hidden), (b, s, self.num_kv_heads, self.head_dim))
@@ -397,6 +409,24 @@ def __init__(
             eps=norm_eps, name="post_per_layer_input_norm"
         )
 
+    def build(self, input_shape):
+        # input_shape = (num_altup_inputs, batch, seq, embed_dim); the active
+        # stream (and every norm/attention/mlp) runs on (batch, seq, embed_dim).
+        active_shape = tuple(input_shape[1:])
+        prefix = active_shape[:-1]
+        self.altup.build(input_shape)
+        self.attention_norm.build(active_shape)
+        self.attention.build(active_shape)
+        self.laurel.build(active_shape)
+        self.post_attention_norm.build(active_shape)
+        self.pre_feedforward_norm.build(active_shape)
+        self.mlp.build(active_shape)
+        self.post_feedforward_norm.build(active_shape)
+        self.per_layer_input_gate.build(active_shape)
+        self.per_layer_projection.build(prefix + (self.hidden_size_per_layer_input,))
+        self.post_per_layer_input_norm.build(active_shape)
+        self.built = True
+
     def finish(self, predictions, active, attn, laurel_output, per_layer_input):
         # Shared post-attention body: AltUp correct + per-layer-input fold-in.
         attn = self.post_attention_norm(attn)
diff --git a/zeromodels/models/glm4v/glm4v_vision_layers.py b/zeromodels/models/glm4v/glm4v_vision_layers.py
index abc2c967..cf5006c8 100644
--- a/zeromodels/models/glm4v/glm4v_vision_layers.py
+++ b/zeromodels/models/glm4v/glm4v_vision_layers.py
@@ -98,6 +98,10 @@ def __init__(self, embed_dim, **kwargs):
         self.embed_dim = embed_dim
         self.proj = layers.Dense(embed_dim, use_bias=True, name="proj")
 
+    def build(self, input_shape):
+        self.proj.build(input_shape)
+        self.built = True
+
     def call(self, x):
         return self.proj(x)
 
@@ -202,6 +206,12 @@ def __init__(self, hidden_size, intermediate_size, **kwargs):
         self.up_proj = layers.Dense(intermediate_size, use_bias=False, name="up")
         self.down_proj = layers.Dense(hidden_size, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate_proj.build(input_shape)
+        self.up_proj.build(input_shape)
+        self.down_proj.build(tuple(input_shape[:-1]) + (self.intermediate_size,))
+        self.built = True
+
     def call(self, x):
         return self.down_proj(ops.silu(self.gate_proj(x)) * self.up_proj(x))
 
@@ -238,6 +248,11 @@ def __init__(self, embed_dim, num_heads, **kwargs):
         self.qkv = layers.Dense(embed_dim * 3, use_bias=False, name="qkv")
         self.proj = layers.Dense(embed_dim, use_bias=False, name="proj")
 
+    def build(self, input_shape):
+        self.qkv.build(input_shape)
+        self.proj.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None):
         seq = ops.shape(hidden_states)[0]
         qkv = self.qkv(hidden_states)
@@ -291,6 +306,13 @@ def __init__(
         self.attn = Glm4vVisionAttention(embed_dim, num_heads, name="attn")
         self.mlp = Glm4VisionMlp(embed_dim, intermediate_size, name="mlp")
 
+    def build(self, input_shape):
+        self.norm1.build(input_shape)
+        self.norm2.build(input_shape)
+        self.attn.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None):
         hidden_states = hidden_states + self.attn(
             self.norm1(hidden_states), cos, sin, attention_mask=attention_mask
@@ -327,6 +349,15 @@ def __init__(self, dim, context_dim, **kwargs):
         self.up_proj = layers.Dense(context_dim, use_bias=False, name="up")
         self.down_proj = layers.Dense(dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.proj.build(input_shape)
+        inter = tuple(input_shape[:-1]) + (self.dim,)
+        self.post_projection_norm.build(inter)
+        self.gate_proj.build(inter)
+        self.up_proj.build(inter)
+        self.down_proj.build(tuple(input_shape[:-1]) + (self.context_dim,))
+        self.built = True
+
     def call(self, x):
         x = self.proj(x)
         x = ops.gelu(self.post_projection_norm(x), approximate=False)
@@ -407,6 +438,21 @@ def __init__(
             out_hidden_size, intermediate_size, name="merger"
         )
 
+    def build(self, input_shape):
+        # input_shape = (num_patches, patch_dim) flattened patches.
+        embed_shape = (None, self.embed_dim)
+        self.patch_embed.build(input_shape)
+        self.post_conv_layernorm.build(embed_shape)
+        self.embeddings.build(embed_shape)
+        for block in self.blocks:
+            block.build(embed_shape)
+        self.post_layernorm.build(embed_shape)
+        self.downsample.build(
+            (None, self.spatial_merge_size, self.spatial_merge_size, self.embed_dim)
+        )
+        self.merger.build((None, self.out_hidden_size))
+        self.built = True
+
     def call(self, pixel_values, grid_thw):
         m = self.spatial_merge_size
         cos, sin = vision_rotary_cos_sin(grid_thw, self.head_dim, m, self.rope_theta)
diff --git a/zeromodels/models/glm5_moe/glm5_moe_layers.py b/zeromodels/models/glm5_moe/glm5_moe_layers.py
index 6076411f..4aec22d5 100644
--- a/zeromodels/models/glm5_moe/glm5_moe_layers.py
+++ b/zeromodels/models/glm5_moe/glm5_moe_layers.py
@@ -65,6 +65,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.silu(self.gate(x)) * self.up(x))
 
@@ -257,6 +263,14 @@ def __init__(
         self.k_norm = layers.LayerNormalization(epsilon=1e-6, name="k_norm")
         self.weights_proj = layers.Dense(n_heads, use_bias=False, name="weights_proj")
 
+    def build(self, input_shape):
+        prefix = tuple(input_shape[:-1])
+        self.wq_b.build(prefix + (self.q_lora_rank,))
+        self.wk.build(input_shape)
+        self.k_norm.build(prefix + (self.head_dim,))
+        self.weights_proj.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, q_resid, cos, sin, attention_mask):
         b = ops.shape(hidden_states)[0]
         s = ops.shape(hidden_states)[1]
@@ -373,6 +387,18 @@ def __init__(
             name="indexer",
         )
 
+    def build(self, input_shape):
+        prefix = tuple(input_shape[:-1])
+        self.q_a_proj.build(input_shape)
+        self.q_a_norm.build(prefix + (self.q_lora_rank,))
+        self.q_b_proj.build(prefix + (self.q_lora_rank,))
+        self.kv_a_proj.build(input_shape)
+        self.kv_a_norm.build(prefix + (self.kv_lora_rank,))
+        self.kv_b_proj.build(prefix + (self.kv_lora_rank,))
+        self.o_proj.build(prefix + (self.num_heads * self.v_head_dim,))
+        self.indexer.build(input_shape)
+        self.built = True
+
     def project_qkv(self, hidden_states, cos, sin):
         b = ops.shape(hidden_states)[0]
         s = ops.shape(hidden_states)[1]
@@ -558,6 +584,13 @@ def __init__(
         else:
             self.mlp = Glm5MoeMLP(embed_dim, mlp_dim, name="mlp")
 
+    def build(self, input_shape):
+        self.attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.mlp_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None, use_cache=False):
         residual = hidden_states
         attn_out = self.attention(
diff --git a/zeromodels/models/gpt_oss/gpt_oss_layers.py b/zeromodels/models/gpt_oss/gpt_oss_layers.py
index a4ec3621..7ed40bae 100644
--- a/zeromodels/models/gpt_oss/gpt_oss_layers.py
+++ b/zeromodels/models/gpt_oss/gpt_oss_layers.py
@@ -144,6 +144,11 @@ def __init__(
         # load time for an mxfp4 checkpoint (the model stays quantization-agnostic).
         self.experts = GptOssExperts(num_experts, embed_dim, mlp_dim, name="experts")
 
+    def build(self, input_shape):
+        self.router.build(input_shape)
+        self.experts.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states):
         b = ops.shape(hidden_states)[0]
         s = ops.shape(hidden_states)[1]
@@ -384,6 +389,13 @@ def __init__(
             name="mlp",
         )
 
+    def build(self, input_shape):
+        self.input_layernorm.build(input_shape)
+        self.self_attn.build(input_shape)
+        self.post_attention_layernorm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/granite_speech5/granite_speech5_layers.py b/zeromodels/models/granite_speech5/granite_speech5_layers.py
index 0c23899f..20173912 100644
--- a/zeromodels/models/granite_speech5/granite_speech5_layers.py
+++ b/zeromodels/models/granite_speech5/granite_speech5_layers.py
@@ -19,6 +19,10 @@ def call(self, attention_mask):
         return ops.cast(ops.all(ops.cast(pairs, "bool"), axis=2), attention_mask.dtype)
 
     def compute_output_shape(self, input_shape):
+        # Some backends (TF) call this with input_shape=None during the symbolic
+        # pass; the halved mask is dynamic either way, so (None, None) is correct.
+        if input_shape is None:
+            return (None, None)
         batch, time = input_shape
         return (batch, None if time is None else time // 2)
 
diff --git a/zeromodels/models/minimax/minimax_layers.py b/zeromodels/models/minimax/minimax_layers.py
index d15361ca..51a2a241 100644
--- a/zeromodels/models/minimax/minimax_layers.py
+++ b/zeromodels/models/minimax/minimax_layers.py
@@ -196,6 +196,15 @@ def __init__(self, embed_dim, num_heads, num_kv_heads, head_dim, **kwargs):
         self.value = layers.Dense(num_kv_heads * head_dim, use_bias=False, name="value")
         self.output_proj = layers.Dense(embed_dim, use_bias=False, name="output_proj")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
+        self.built = True
+
     def project_qkv(self, hidden_states):
         b = ops.shape(hidden_states)[0]
         s = ops.shape(hidden_states)[1]
@@ -347,6 +356,14 @@ def __init__(
         # independent of config.rms_norm_eps.
         self.norm = MiniMaxRMSNorm(eps=1e-6, name="norm")
 
+    def build(self, input_shape):
+        proj_shape = tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        self.qkv.build(input_shape)
+        self.output_gate.build(input_shape)
+        self.norm.build(proj_shape)
+        self.output_proj.build(proj_shape)
+        self.built = True
+
     def split_qkv(self, hidden_states):
         b = ops.shape(hidden_states)[0]
         s = ops.shape(hidden_states)[1]
@@ -516,6 +533,13 @@ def __init__(
             num_experts, num_experts_per_tok, embed_dim, mlp_dim, name="mlp"
         )
 
+    def build(self, input_shape):
+        self.attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.mlp_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def run_mlp(self, hidden_states):
         h = self.mlp_norm(hidden_states)
         return h * self.mlp_alpha + self.mlp(h) * self.mlp_beta
diff --git a/zeromodels/models/minimax_m3_vl/minimax_m3_vl_layers.py b/zeromodels/models/minimax_m3_vl/minimax_m3_vl_layers.py
index 275e36ba..4db804ae 100644
--- a/zeromodels/models/minimax_m3_vl/minimax_m3_vl_layers.py
+++ b/zeromodels/models/minimax_m3_vl/minimax_m3_vl_layers.py
@@ -699,6 +699,13 @@ def __init__(self, embed_dim, num_heads, **kwargs):
         self.value = layers.Dense(embed_dim, name="value")
         self.output_proj = layers.Dense(embed_dim, name="output_proj")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin):
         b = ops.shape(hidden_states)[0]
         s = ops.shape(hidden_states)[1]
@@ -755,6 +762,14 @@ def __init__(self, embed_dim, mlp_dim, num_heads, norm_eps=1e-5, **kwargs):
         self.fc1 = layers.Dense(mlp_dim, name="fc1")
         self.fc2 = layers.Dense(embed_dim, name="fc2")
 
+    def build(self, input_shape):
+        self.layer_norm1.build(input_shape)
+        self.attention.build(input_shape)
+        self.layer_norm2.build(input_shape)
+        self.fc1.build(input_shape)
+        self.fc2.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, hidden_states, cos, sin):
         hidden_states = hidden_states + self.attention(
             self.layer_norm1(hidden_states), cos, sin
diff --git a/zeromodels/models/minimax_m3_vl/minimax_m3_vl_model.py b/zeromodels/models/minimax_m3_vl/minimax_m3_vl_model.py
index d81ed6cb..40e630bb 100644
--- a/zeromodels/models/minimax_m3_vl/minimax_m3_vl_model.py
+++ b/zeromodels/models/minimax_m3_vl/minimax_m3_vl_model.py
@@ -75,6 +75,14 @@ def __init__(
             for i in range(num_layers)
         ]
 
+    def build(self, input_shape):
+        # input_shape = (num_patches, patch_dim) packed patches.
+        self.patch_embed.build(input_shape)
+        self.pre_norm.build((None, self.embed_dim))
+        for block in self.blocks:
+            block.build((None, None, self.embed_dim))
+        self.built = True
+
     def call(self, pixel_values, grid_thw=None):
         # grid_thw drives host-side rope construction (``for t,h,w in grid_thw``),
         # so it must be a concrete iterable. In the functional graph this call runs
diff --git a/zeromodels/models/qwen2/qwen2_layers.py b/zeromodels/models/qwen2/qwen2_layers.py
index 1a93516f..85471c76 100644
--- a/zeromodels/models/qwen2/qwen2_layers.py
+++ b/zeromodels/models/qwen2/qwen2_layers.py
@@ -62,6 +62,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.silu(self.gate(x)) * self.up(x))
 
@@ -119,6 +125,15 @@ def __init__(self, embed_dim, num_heads, num_kv_heads, head_dim=None, **kwargs):
         )
         self.output_proj = layers.Dense(embed_dim, use_bias=False, name="output_proj")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
+        self.built = True
+
     def _split_heads(self, x, num_heads):
         b = ops.shape(x)[0]
         s = ops.shape(x)[1]
@@ -261,6 +276,13 @@ def __init__(
         self.mlp_norm = Qwen2RMSNorm(eps=norm_eps, name="mlp_norm")
         self.mlp = Qwen2MLP(embed_dim, mlp_dim, name="mlp")
 
+    def build(self, input_shape):
+        self.attention_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.mlp_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/qwen2_5_vl/qwen2_5_vl_layers.py b/zeromodels/models/qwen2_5_vl/qwen2_5_vl_layers.py
index 45033e01..63faa3d4 100644
--- a/zeromodels/models/qwen2_5_vl/qwen2_5_vl_layers.py
+++ b/zeromodels/models/qwen2_5_vl/qwen2_5_vl_layers.py
@@ -65,6 +65,12 @@ def __init__(self, embed_dim, mlp_dim, use_bias=False, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=use_bias, name="up")
         self.down = layers.Dense(embed_dim, use_bias=use_bias, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.silu(self.gate(x)) * self.up(x))
 
@@ -359,6 +365,10 @@ def __init__(self, embed_dim, **kwargs):
         self.embed_dim = embed_dim
         self.proj = layers.Dense(embed_dim, use_bias=False, name="proj")
 
+    def build(self, input_shape):
+        self.proj.build(input_shape)
+        self.built = True
+
     def call(self, x):
         return self.proj(x)
 
@@ -397,6 +407,11 @@ def __init__(self, embed_dim, num_heads, **kwargs):
         self.qkv = layers.Dense(embed_dim * 3, use_bias=True, name="qkv")
         self.proj = layers.Dense(embed_dim, use_bias=True, name="proj")
 
+    def build(self, input_shape):
+        self.qkv.build(input_shape)
+        self.proj.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None):
         seq = ops.shape(hidden_states)[0]
         qkv = ops.reshape(
@@ -459,6 +474,13 @@ def __init__(self, embed_dim, num_heads, intermediate_size, **kwargs):
         self.attn = Qwen2_5VLVisionAttention(embed_dim, num_heads, name="attn")
         self.mlp = Qwen2_5VLMLP(embed_dim, intermediate_size, use_bias=True, name="mlp")
 
+    def build(self, input_shape):
+        self.norm1.build(input_shape)
+        self.norm2.build(input_shape)
+        self.attn.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None):
         hidden_states = hidden_states + self.attn(
             self.norm1(hidden_states), cos, sin, attention_mask=attention_mask
@@ -504,6 +526,13 @@ def __init__(self, dim, context_dim, spatial_merge_size=2, **kwargs):
         self.mlp_fc1 = layers.Dense(self.hidden_size, use_bias=True, name="mlp_fc1")
         self.mlp_fc2 = layers.Dense(dim, use_bias=True, name="mlp_fc2")
 
+    def build(self, input_shape):
+        self.ln_q.build(input_shape)
+        merged_shape = (None, self.hidden_size)
+        self.mlp_fc1.build(merged_shape)
+        self.mlp_fc2.build(merged_shape)
+        self.built = True
+
     def call(self, x):
         x = ops.reshape(self.ln_q(x), (-1, self.hidden_size))
         return self.mlp_fc2(ops.gelu(self.mlp_fc1(x), approximate=False))
diff --git a/zeromodels/models/qwen2_5_vl/qwen2_5_vl_model.py b/zeromodels/models/qwen2_5_vl/qwen2_5_vl_model.py
index 42fa7629..4f0c4a78 100644
--- a/zeromodels/models/qwen2_5_vl/qwen2_5_vl_model.py
+++ b/zeromodels/models/qwen2_5_vl/qwen2_5_vl_model.py
@@ -160,6 +160,15 @@ def __init__(
             name="merger",
         )
 
+    def build(self, input_shape):
+        # input_shape = (num_patches, patch_dim) packed patches.
+        embed_shape = (None, self.embed_dim)
+        self.patch_embed.build(input_shape)
+        for block in self.blocks:
+            block.build(embed_shape)
+        self.merger.build(embed_shape)
+        self.built = True
+
     def call(self, pixel_values, grid_thw):
         grid_rows = [
             tuple(int(v) for v in row)
diff --git a/zeromodels/models/qwen2_vl/qwen2_vl_layers.py b/zeromodels/models/qwen2_vl/qwen2_vl_layers.py
index 571b97e9..c7322472 100644
--- a/zeromodels/models/qwen2_vl/qwen2_vl_layers.py
+++ b/zeromodels/models/qwen2_vl/qwen2_vl_layers.py
@@ -329,6 +329,10 @@ def __init__(self, embed_dim, use_bias=False, **kwargs):
         self.use_bias = use_bias
         self.proj = layers.Dense(embed_dim, use_bias=use_bias, name="proj")
 
+    def build(self, input_shape):
+        self.proj.build(input_shape)
+        self.built = True
+
     def call(self, x):
         return self.proj(x)
 
@@ -362,6 +366,11 @@ def __init__(self, embed_dim, num_heads, **kwargs):
         self.qkv = layers.Dense(embed_dim * 3, use_bias=True, name="qkv")
         self.proj = layers.Dense(embed_dim, use_bias=True, name="proj")
 
+    def build(self, input_shape):
+        self.qkv.build(input_shape)
+        self.proj.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None):
         seq = ops.shape(hidden_states)[0]
         qkv = self.qkv(hidden_states)
@@ -405,6 +414,11 @@ def __init__(self, embed_dim, hidden_dim, **kwargs):
         self.fc1 = layers.Dense(hidden_dim, use_bias=True, name="fc1")
         self.fc2 = layers.Dense(embed_dim, use_bias=True, name="fc2")
 
+    def build(self, input_shape):
+        self.fc1.build(input_shape)
+        self.fc2.build(tuple(input_shape[:-1]) + (self.hidden_dim,))
+        self.built = True
+
     def call(self, x):
         return self.fc2(quick_gelu(self.fc1(x)))
 
@@ -432,6 +446,13 @@ def __init__(self, embed_dim, num_heads, mlp_ratio=4, **kwargs):
         self.attn = Qwen2VLVisionAttention(embed_dim, num_heads, name="attn")
         self.mlp = Qwen2VLVisionMLP(embed_dim, int(embed_dim * mlp_ratio), name="mlp")
 
+    def build(self, input_shape):
+        self.norm1.build(input_shape)
+        self.norm2.build(input_shape)
+        self.attn.build(input_shape)
+        self.mlp.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, cos, sin, attention_mask=None):
         hidden_states = hidden_states + self.attn(
             self.norm1(hidden_states), cos, sin, attention_mask=attention_mask
@@ -477,6 +498,13 @@ def __init__(
         self.mlp_fc1 = layers.Dense(self.hidden_size, use_bias=True, name="mlp_fc1")
         self.mlp_fc2 = layers.Dense(dim, use_bias=True, name="mlp_fc2")
 
+    def build(self, input_shape):
+        self.ln_q.build(input_shape)
+        merged_shape = (None, self.hidden_size)
+        self.mlp_fc1.build(merged_shape)
+        self.mlp_fc2.build(merged_shape)
+        self.built = True
+
     def call(self, x):
         x = self.ln_q(x)
         x = ops.reshape(x, (-1, self.hidden_size))
diff --git a/zeromodels/models/qwen2_vl/qwen2_vl_model.py b/zeromodels/models/qwen2_vl/qwen2_vl_model.py
index 5411af0b..81d0e098 100644
--- a/zeromodels/models/qwen2_vl/qwen2_vl_model.py
+++ b/zeromodels/models/qwen2_vl/qwen2_vl_model.py
@@ -241,6 +241,15 @@ def __init__(
             llm_hidden_size, embed_dim, spatial_merge_size, name="merger"
         )
 
+    def build(self, input_shape):
+        # input_shape = (num_patches, patch_dim) packed patches.
+        embed_shape = (None, self.embed_dim)
+        self.patch_embed.build(input_shape)
+        for block in self.blocks:
+            block.build(embed_shape)
+        self.merger.build(embed_shape)
+        self.built = True
+
     def call(self, pixel_values, grid_thw):
         cos, sin = vision_rotary_cos_sin(
             grid_thw, self.head_dim, self.spatial_merge_size
diff --git a/zeromodels/models/qwen3_5/qwen3_5_layers.py b/zeromodels/models/qwen3_5/qwen3_5_layers.py
index 62862448..2bfed7fb 100644
--- a/zeromodels/models/qwen3_5/qwen3_5_layers.py
+++ b/zeromodels/models/qwen3_5/qwen3_5_layers.py
@@ -104,6 +104,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.silu(self.gate(x)) * self.up(x))
 
@@ -174,6 +180,15 @@ def __init__(
         self.query_norm = Qwen3_5RMSNorm(eps=norm_eps, name="query_norm")
         self.key_norm = Qwen3_5RMSNorm(eps=norm_eps, name="key_norm")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.query_norm.build(tuple(input_shape[:-1]) + (self.num_heads, self.head_dim))
+        self.key_norm.build(tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim))
+        self.built = True
+
     def call(
         self,
         hidden_states,
@@ -559,6 +574,16 @@ def __init__(self, config, layer_type, **kwargs):
                 name="linear_attn",
             )
 
+    def build(self, input_shape):
+        self.attention_norm.build(input_shape)
+        self.mlp_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        if self.layer_type == "full_attention":
+            self.attention.build(input_shape)
+        else:
+            self.linear_attn.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/qwen3_5_moe/qwen3_5_moe_model.py b/zeromodels/models/qwen3_5_moe/qwen3_5_moe_model.py
index 1f68b784..06893da9 100644
--- a/zeromodels/models/qwen3_5_moe/qwen3_5_moe_model.py
+++ b/zeromodels/models/qwen3_5_moe/qwen3_5_moe_model.py
@@ -300,6 +300,13 @@ def __init__(
         ]
         self.final_norm = Qwen3NextRMSNorm(eps=norm_eps, name="final_norm")
 
+    def build(self, input_shape):
+        self.token_embedding.build((input_shape[0], input_shape[1]))
+        for layer in self.decoder_layers:
+            layer.build(input_shape)
+        self.final_norm.build(input_shape)
+        self.built = True
+
     def call(
         self,
         inputs_embeds,
diff --git a/zeromodels/models/qwen3_next/qwen3_next_layers.py b/zeromodels/models/qwen3_next/qwen3_next_layers.py
index 84ab9c16..ab5668b5 100644
--- a/zeromodels/models/qwen3_next/qwen3_next_layers.py
+++ b/zeromodels/models/qwen3_next/qwen3_next_layers.py
@@ -104,6 +104,12 @@ def __init__(self, embed_dim, mlp_dim, **kwargs):
         self.up = layers.Dense(mlp_dim, use_bias=False, name="up")
         self.down = layers.Dense(embed_dim, use_bias=False, name="down")
 
+    def build(self, input_shape):
+        self.gate.build(input_shape)
+        self.up.build(input_shape)
+        self.down.build(tuple(input_shape[:-1]) + (self.mlp_dim,))
+        self.built = True
+
     def call(self, x):
         return self.down(ops.silu(self.gate(x)) * self.up(x))
 
@@ -174,6 +180,19 @@ def __init__(
         self.query_norm = Qwen3NextRMSNorm(eps=norm_eps, name="query_norm")
         self.key_norm = Qwen3NextRMSNorm(eps=norm_eps, name="key_norm")
 
+    def build(self, input_shape):
+        self.query.build(input_shape)
+        self.key.build(input_shape)
+        self.value.build(input_shape)
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
+        self.query_norm.build(tuple(input_shape[:-1]) + (self.num_heads, self.head_dim))
+        self.key_norm.build(
+            tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim)
+        )
+        self.built = True
+
     def call(
         self,
         hidden_states,
@@ -712,6 +731,16 @@ def __init__(self, config, layer_type, use_moe=True, **kwargs):
                 name="linear_attn",
             )
 
+    def build(self, input_shape):
+        self.attention_norm.build(input_shape)
+        self.mlp_norm.build(input_shape)
+        self.mlp.build(input_shape)
+        if self.layer_type == "full_attention":
+            self.attention.build(input_shape)
+        else:
+            self.linear_attn.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,
diff --git a/zeromodels/models/t5/t5_layers.py b/zeromodels/models/t5/t5_layers.py
index e28a761b..f75fa730 100644
--- a/zeromodels/models/t5/t5_layers.py
+++ b/zeromodels/models/t5/t5_layers.py
@@ -93,6 +93,13 @@ def __init__(self, embed_dim, key_value_dim, num_heads, prefix, **kwargs):
         self.v = layers.Dense(self.inner_dim, use_bias=False, name=f"{prefix}_v")
         self.o = layers.Dense(embed_dim, use_bias=False, name=f"{prefix}_o")
 
+    def build(self, input_shape):
+        self.q.build(input_shape)
+        self.k.build(input_shape)
+        self.v.build(input_shape)
+        self.o.build(tuple(input_shape[:-1]) + (self.inner_dim,))
+        self.built = True
+
     def split_heads(self, x, batch):
         x = ops.reshape(x, (batch, -1, self.num_heads, self.key_value_dim))
         return ops.transpose(x, (0, 2, 1, 3))
@@ -121,6 +128,11 @@ def __init__(self, embed_dim, key_value_dim, num_heads, eps, prefix, **kwargs):
         self.layer_norm = T5LayerNorm(eps, name=f"{prefix}_ln")
         self.attention = T5Attention(embed_dim, key_value_dim, num_heads, prefix=prefix)
 
+    def build(self, input_shape):
+        self.layer_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, position_bias):
         normed = self.layer_norm(hidden_states)
         return hidden_states + self.attention(normed, position_bias)
@@ -132,6 +144,11 @@ def __init__(self, embed_dim, key_value_dim, num_heads, eps, prefix, **kwargs):
         self.layer_norm = T5LayerNorm(eps, name=f"{prefix}_ln")
         self.attention = T5Attention(embed_dim, key_value_dim, num_heads, prefix=prefix)
 
+    def build(self, input_shape):
+        self.layer_norm.build(input_shape)
+        self.attention.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, encoder_hidden_states, position_bias):
         normed = self.layer_norm(hidden_states)
         attn = self.attention(
@@ -148,6 +165,14 @@ def __init__(self, embed_dim, mlp_dim, hidden_act, eps, prefix, **kwargs):
         self.act = layers.Activation(hidden_act, name=f"{prefix}_act")
         self.wo = layers.Dense(embed_dim, use_bias=False, name=f"{prefix}_wo")
 
+    def build(self, input_shape):
+        self.layer_norm.build(input_shape)
+        self.wi.build(input_shape)
+        inter_shape = tuple(input_shape[:-1]) + (self.wi.units,)
+        self.act.build(inter_shape)
+        self.wo.build(inter_shape)
+        self.built = True
+
     def call(self, hidden_states):
         normed = self.layer_norm(hidden_states)
         forwarded = self.wo(self.act(self.wi(normed)))
@@ -174,6 +199,11 @@ def __init__(
             embed_dim, mlp_dim, hidden_act, eps, prefix=f"{prefix}_ff"
         )
 
+    def build(self, input_shape):
+        self.self_attention.build(input_shape)
+        self.ff.build(input_shape)
+        self.built = True
+
     def call(self, hidden_states, position_bias):
         hidden_states = self.self_attention(hidden_states, position_bias)
         return self.ff(hidden_states)
@@ -205,6 +235,12 @@ def __init__(
             embed_dim, mlp_dim, hidden_act, eps, prefix=f"{prefix}_ff"
         )
 
+    def build(self, input_shape):
+        self.self_attention.build(input_shape)
+        self.cross_attention.build(input_shape)
+        self.ff.build(input_shape)
+        self.built = True
+
     def call(
         self,
         hidden_states,

From c7590369628d37abd162860260725ec703e96ec5 Mon Sep 17 00:00:00 2001
From: IMvision12 <gitesh.ch.0912@gmail.com>
Date: Fri, 28 Aug 2026 16:05:46 -0700
Subject: [PATCH 2/3] Format

---
 zeromodels/models/gemma/gemma_layers.py     | 4 +++-
 zeromodels/models/gemma2/gemma2_layers.py   | 4 +++-
 zeromodels/models/gemma3/gemma3_layers.py   | 8 ++++++--
 zeromodels/models/qwen3_5/qwen3_5_layers.py | 8 ++++++--
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/zeromodels/models/gemma/gemma_layers.py b/zeromodels/models/gemma/gemma_layers.py
index 22904ac0..49398ee1 100644
--- a/zeromodels/models/gemma/gemma_layers.py
+++ b/zeromodels/models/gemma/gemma_layers.py
@@ -122,7 +122,9 @@ def build(self, input_shape):
         self.query.build(input_shape)
         self.key.build(input_shape)
         self.value.build(input_shape)
-        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
         self.built = True
 
     def call(
diff --git a/zeromodels/models/gemma2/gemma2_layers.py b/zeromodels/models/gemma2/gemma2_layers.py
index af756c17..6fac518a 100644
--- a/zeromodels/models/gemma2/gemma2_layers.py
+++ b/zeromodels/models/gemma2/gemma2_layers.py
@@ -129,7 +129,9 @@ def build(self, input_shape):
         self.query.build(input_shape)
         self.key.build(input_shape)
         self.value.build(input_shape)
-        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
         self.built = True
 
     def softcap(self, attn):
diff --git a/zeromodels/models/gemma3/gemma3_layers.py b/zeromodels/models/gemma3/gemma3_layers.py
index 5d40d1a9..2eb8f26f 100644
--- a/zeromodels/models/gemma3/gemma3_layers.py
+++ b/zeromodels/models/gemma3/gemma3_layers.py
@@ -134,9 +134,13 @@ def build(self, input_shape):
         self.query.build(input_shape)
         self.key.build(input_shape)
         self.value.build(input_shape)
-        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
         self.query_norm.build(tuple(input_shape[:-1]) + (self.num_heads, self.head_dim))
-        self.key_norm.build(tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim))
+        self.key_norm.build(
+            tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim)
+        )
         self.built = True
 
     def call(
diff --git a/zeromodels/models/qwen3_5/qwen3_5_layers.py b/zeromodels/models/qwen3_5/qwen3_5_layers.py
index 2bfed7fb..a1c0edc3 100644
--- a/zeromodels/models/qwen3_5/qwen3_5_layers.py
+++ b/zeromodels/models/qwen3_5/qwen3_5_layers.py
@@ -184,9 +184,13 @@ def build(self, input_shape):
         self.query.build(input_shape)
         self.key.build(input_shape)
         self.value.build(input_shape)
-        self.output_proj.build(tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,))
+        self.output_proj.build(
+            tuple(input_shape[:-1]) + (self.num_heads * self.head_dim,)
+        )
         self.query_norm.build(tuple(input_shape[:-1]) + (self.num_heads, self.head_dim))
-        self.key_norm.build(tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim))
+        self.key_norm.build(
+            tuple(input_shape[:-1]) + (self.num_kv_heads, self.head_dim)
+        )
         self.built = True
 
     def call(

From b300477e8c662510ce663dc7b514be3304f84c06 Mon Sep 17 00:00:00 2001
From: IMvision12 <gitesh.ch.0912@gmail.com>
Date: Fri, 28 Aug 2026 16:13:00 -0700
Subject: [PATCH 3/3] Fix pre commit

---
 .pre-commit-config.yaml | 2 +-
 docs/efficientdet.md    | 2 +-
 pyproject.toml          | 8 ++++++++
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1205220f..578c75b5 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -11,7 +11,7 @@ repos:
         args: ["--maxkb=2000"]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.10
+    rev: v0.16.1
     hooks:
       - id: ruff
         args: [--fix]
diff --git a/docs/efficientdet.md b/docs/efficientdet.md
index f206e837..f9081250 100644
--- a/docs/efficientdet.md
+++ b/docs/efficientdet.md
@@ -398,4 +398,4 @@ output = model(inputs["pixel_values"], training=False)
 
 Set it once at the top of a script, since already-built models keep the layout they were
 constructed with. The post-processor emits `xyxy` pixel boxes and class indices, which have no
-channel axis, so it is not format-sensitive.
\ No newline at end of file
+channel axis, so it is not format-sensitive.
diff --git a/pyproject.toml b/pyproject.toml
index 18308f4d..2f46cbdd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -72,6 +72,14 @@ conversion = [
     "timm",
     "transformers",
 ]
+# Developer tooling: the pre-commit hooks (ruff lint + format, whitespace/EOF
+# fixers). Install with: pip install -e .[dev], then run: pre-commit install
+# Keep ruff pinned to the .pre-commit-config.yaml rev so a manual `ruff` run
+# matches what the hook enforces.
+dev = [
+    "pre-commit",
+    "ruff==0.16.1",
+]
 
 [project.scripts]
 zeromodels-test = "tests._test_runner:main"