Fix TF SparseTensor issue and pre commit fix

.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]

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.
+channel axis, so it is not format-sensitive.

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"

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",

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)

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,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 call(
         self,
         hidden_states,
@@ -250,6 +265,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,

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,15 @@ 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 +305,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,

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,19 @@ 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 +309,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,

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)

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)