feat(vla): add SmolVLA conditioning and experiment artifacts

tests/test_imf_transformer1d_external_alignment.py

@@ -191,6 +191,73 @@ class IMFTransformer1DExternalAlignmentTest(unittest.TestCase):
                 self._optim_group_names(external_model, external_groups),
             )
 
+    def test_attnres_full_supports_multihead_grouped_query_attention(self):
+        local_module = _load_local_module()
+        model = local_module.IMFTransformer1D(
+            input_dim=4,
+            output_dim=4,
+            horizon=6,
+            n_obs_steps=3,
+            cond_dim=5,
+            n_layer=1,
+            n_head=4,
+            n_kv_head=2,
+            n_emb=16,
+            p_drop_emb=0.0,
+            p_drop_attn=0.0,
+            causal_attn=False,
+            time_as_cond=True,
+            n_cond_layers=0,
+            backbone_type='attnres_full',
+        )
+        model.eval()
+
+        attention = model.attnres_backbone.layers[0].fn
+        self.assertEqual(attention.n_heads, 4)
+        self.assertEqual(attention.n_kv_heads, 2)
+        self.assertEqual(attention.d_head, 4)
+
+        sample = torch.randn(2, 6, 4)
+        r = torch.tensor([0.1, 0.4], dtype=torch.float32)
+        t = torch.tensor([0.7, 0.9], dtype=torch.float32)
+        cond = torch.randn(2, 3, 5)
+        with torch.no_grad():
+            output = model(sample=sample, r=r, t=t, cond=cond)
+
+        self.assertEqual(output.shape, (2, 6, 4))
+
+    def test_attnres_full_rejects_invalid_multihead_shapes(self):
+        local_module = _load_local_module()
+        with self.assertRaisesRegex(ValueError, 'must be divisible by n_head'):
+            local_module.IMFTransformer1D(
+                input_dim=4,
+                output_dim=4,
+                horizon=6,
+                n_head=5,
+                n_emb=16,
+                backbone_type='attnres_full',
+            )
+        with self.assertRaisesRegex(ValueError, 'n_head=4 must be divisible by n_kv_head=3'):
+            local_module.IMFTransformer1D(
+                input_dim=4,
+                output_dim=4,
+                horizon=6,
+                n_head=4,
+                n_kv_head=3,
+                n_emb=16,
+                backbone_type='attnres_full',
+            )
+        with self.assertRaisesRegex(ValueError, 'requires an even per-head dimension'):
+            local_module.IMFTransformer1D(
+                input_dim=4,
+                output_dim=4,
+                horizon=6,
+                n_head=2,
+                n_kv_head=1,
+                n_emb=18,
+                backbone_type='attnres_full',
+            )
+
 
 if __name__ == '__main__':
     unittest.main()