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

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

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

{broccoli_ml-13.0.5 → 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.5 → broccoli_ml-13.0.6}/pyproject.toml

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