broccoli-ml 13.0.4__tar.gz → 13.0.6__tar.gz

{broccoli_ml-13.0.4 → broccoli_ml-13.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: broccoli-ml
-Version: 13.0.4
+Version: 13.0.6
 Summary: Some useful Pytorch models, circa 2025
 License: MIT
 Author: Nicholas Bailey

{broccoli_ml-13.0.4 → broccoli_ml-13.0.6}/broccoli/transformer.py

@@ -200,26 +200,26 @@ class MHAttention(nn.Module):
                 "`source_size` must be a tuple of 1, 2 or 3 integers"
             )
 
-        q = rearrange(q, "b t (h d) -> b t h d", h=self.n_heads)
-        k = rearrange(k, "b t (h d) -> b t h d", h=self.n_heads)
+        q = rearrange(q, "b t (h d) -> b h t d", h=self.n_heads)
+        k = rearrange(k, "b t (h d) -> b h t d", h=self.n_heads)
 
         q_util, q_img = (
-            q[:, : self.utility_tokens, :, :],
-            q[:, self.utility_tokens :, :, :],
+            q[:, :, : self.utility_tokens, :],
+            q[:, :, self.utility_tokens :, :],
         )
         k_util, k_img = (
-            k[:, : self.utility_tokens, :, :],
-            k[:, self.utility_tokens :, :, :],
+            k[:, :, : self.utility_tokens, :],
+            k[:, :, self.utility_tokens :, :],
         )
 
         q_img = rearrange(
             q_img,
-            f"b ({spatial_dimension_names}) h d -> b {spatial_dimension_names} h d",
+            f"b h ({spatial_dimension_names}) d -> b h {spatial_dimension_names} d",
             **spatial_dimension_values,
         )
         k_img = rearrange(
             k_img,
-            f"b ({spatial_dimension_names}) h d -> b {spatial_dimension_names} h d",
+            f"b h ({spatial_dimension_names}) d -> b h {spatial_dimension_names} d",
             **spatial_dimension_values,
         )
 
@@ -230,19 +230,19 @@ class MHAttention(nn.Module):
 
         q_img = rearrange(
             q_img,
-            f"b {spatial_dimension_names} h d -> b ({spatial_dimension_names}) h d",
+            f"b h {spatial_dimension_names} d -> b h ({spatial_dimension_names}) d",
         )
         k_img = rearrange(
             k_img,
-            f"b {spatial_dimension_names} h d -> b ({spatial_dimension_names}) h d",
+            f"b h {spatial_dimension_names} d -> b h ({spatial_dimension_names}) d",
         )
 
         # Re-combine the utility tokens and the RoPE-enhanced sequence tokens
-        q = torch.cat([q_util, q_img], dim=1)
-        k = torch.cat([k_util, k_img], dim=1)
+        q = torch.cat([q_util, q_img], dim=2)
+        k = torch.cat([k_util, k_img], dim=2)
 
-        q = rearrange(q, "b t h d -> b t (h d)")
-        k = rearrange(k, "b t h d -> b t (h d)")
+        q = rearrange(q, "b h t d -> b t (h d)")
+        k = rearrange(k, "b h t d -> b t (h d)")
 
         return q, k
 

{broccoli_ml-13.0.4 → broccoli_ml-13.0.6}/broccoli/vit.py

@@ -433,9 +433,11 @@ class ViTEncoder(nn.Module):
         return self.transformer.attention_logits(x)
 
     def reset_parameters(self):
-        for module in self.encoder:
+        for module in self.preprocess:
             if hasattr(module, "reset_parameters"):
                 module.reset_parameters()
+        self.initial_ff.reset_parameters()
+        self.transformer.reset_parameters()
 
 
 class ViT(nn.Module):

{broccoli_ml-13.0.4 → broccoli_ml-13.0.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "broccoli-ml"
-version = "13.0.4"
+version = "13.0.6"
 description = "Some useful Pytorch models, circa 2025"
 authors = [
     {name = "Nicholas Bailey"}