[BugFix] pre-commit

Author: RandallBalestriero

stable_pretraining/backbone/video/recurrent_vit.py

@@ -16,7 +16,8 @@
   spatial stack uses an additive learned pos embed.
 - Per-frame CLS token prepended; the full ``(1+P, D)`` sequence goes through
   ``spatial_depth`` :class:`~stable_pretraining.backbone.vit.TransformerBlock`
-  layers (``self_attn=True, use_rope=None, use_qk_norm=True, mlp_type='gelu'``).
+  layers (``self_attn=True, use_rope=None, mlp_type='gelu'``). QK-norm is
+  off by default and only enabled on the ``huge`` preset.
 - ``LayerNorm`` (``norm_s``) then split: CLS → ``(B, T, D)``, patches →
   ``(B, T, P, D)``.
 - ``nn.GRU(D, D, num_layers=gru_layers, batch_first=True)`` walks the CLS
@@ -47,8 +48,8 @@
     x = torch.randn(2, 3, 8, 64, 64)
     out = enc(x)
     out.feature_map.shape  # (2, 192, 8, 8, 8)
-    out.pooled.shape       # (2, 192)
-    out.tokens.shape       # (2, 8, 192)
+    out.pooled.shape  # (2, 192)
+    out.tokens.shape  # (2, 8, 192)
 """
 
 from __future__ import annotations
@@ -244,8 +245,10 @@ def forward(
             # (B, 1, D) -> broadcast to (B*T, 1, D) so each frame of clip i
             # carries clip i's cam embedding.
             ce = self.cam_embed[cam_id]
-            ce = ce.unsqueeze(1).expand(B, T, 1, self.embed_dim).reshape(
-                B * T, 1, self.embed_dim
+            ce = (
+                ce.unsqueeze(1)
+                .expand(B, T, 1, self.embed_dim)
+                .reshape(B * T, 1, self.embed_dim)
             )
             x = x + ce
 
@@ -270,13 +273,9 @@ def forward(
             .permute(0, 4, 1, 2, 3)
             .contiguous()
         )
-        pooled = (
-            feature_map.mean(dim=(2, 3, 4)) if self.global_pool == "avg" else None
-        )
+        pooled = feature_map.mean(dim=(2, 3, 4)) if self.global_pool == "avg" else None
 
-        return RecurrentViTOutput(
-            feature_map=feature_map, pooled=pooled, tokens=tokens
-        )
+        return RecurrentViTOutput(feature_map=feature_map, pooled=pooled, tokens=tokens)
 
 
 # -----------------------------------------------------------------------------

stable_pretraining/tests/unit/test_recurrent_vit.py

@@ -22,8 +22,11 @@ class TestRecurrentViTShapes:
     """Output shape and dataclass contract."""
 
     def test_recurrent_vit_shapes(self):
-        """Default tiny on a (2, 3, 8, 64, 64) clip produces the documented
-        ``feature_map`` / ``pooled`` / ``tokens`` shapes."""
+        """Default tiny produces the documented output shapes.
+
+        Input ``(2, 3, 8, 64, 64)`` → ``feature_map=(2, 192, 8, 8, 8)``,
+        ``pooled=(2, 192)``, ``tokens=(2, 8, 192)``.
+        """
         torch.manual_seed(0)
         enc = recurrent_vit_tiny()
         enc.eval()
@@ -37,8 +40,7 @@ def test_recurrent_vit_shapes(self):
         assert out.tokens.shape == (2, 8, 192)
 
     def test_no_pool(self):
-        """``global_pool=''`` zeroes the pooled slot but keeps feature_map and
-        tokens populated."""
+        """``global_pool=''`` drops the pooled slot but keeps feature_map and tokens."""
         m = RecurrentViT(
             img_size=32,
             patch_size=8,
@@ -70,9 +72,12 @@ def test_rejects_wrong_spatial_size(self):
 
 @pytest.mark.unit
 class TestNoFutureLeakage:
-    """The GRU on the pooled CLS sequence makes the model causal in time;
-    the per-frame ViT can't leak future info either because frames are
-    encoded independently."""
+    """Causality of feature_map and tokens.
+
+    The GRU on the pooled CLS sequence makes ``tokens`` causal in time;
+    the per-frame ViT can't leak future info into ``feature_map`` either,
+    because frames are encoded independently.
+    """
 
     def test_no_future_leakage(self):
         torch.manual_seed(0)
@@ -92,21 +97,19 @@ def test_no_future_leakage(self):
             out_a.feature_map[:, :, 4:], out_b.feature_map[:, :, 4:], atol=1e-5
         )
         # tokens are GRU output: causal by construction.
-        assert torch.allclose(
-            out_a.tokens[:, :4], out_b.tokens[:, :4], atol=1e-5
-        )
-        assert not torch.allclose(
-            out_a.tokens[:, 4:], out_b.tokens[:, 4:], atol=1e-5
-        )
+        assert torch.allclose(out_a.tokens[:, :4], out_b.tokens[:, :4], atol=1e-5)
+        assert not torch.allclose(out_a.tokens[:, 4:], out_b.tokens[:, 4:], atol=1e-5)
 
 
 @pytest.mark.unit
 class TestFactoryParamCounts:
-    """Each preset's total parameter count must be in the documented
-    ViT-family budget ballpark. Tolerance is ±20%: the targets are
-    standard-ViT-shaped approximations, and the GRU + patch embed + QK-norm
-    each shift the count by a few percent, so a single fixed target won't
-    sit at the center of every preset."""
+    """Each preset lands in the documented ViT-family param budget.
+
+    Tolerance is ±20%: the targets are standard-ViT-shaped approximations,
+    and the GRU + patch embed + QK-norm each shift the count by a few
+    percent, so a single fixed target won't sit at the center of every
+    preset.
+    """
 
     @pytest.mark.parametrize(
         "factory,target",
@@ -133,8 +136,11 @@ def test_factory_param_counts(self, factory, target):
 
 @pytest.mark.unit
 class TestSmallerPresetsForward:
-    """Smoke-forward the smaller presets so we cover the GRU path on a real
-    tensor without paying the large preset's memory cost."""
+    """Real-tensor smoke forward for the smaller presets.
+
+    Exercises the GRU path on a real tensor without paying the large
+    preset's memory cost.
+    """
 
     @pytest.mark.parametrize("factory", [recurrent_vit_tiny, recurrent_vit_small])
     def test_forward_runs(self, factory):