broccoli-ml 5.1.2__py3-none-any.whl → 5.1.4__py3-none-any.whl

broccoli/transformer.py

@@ -284,7 +284,7 @@ class MHAttention(nn.Module):
 
             return self.out_proj(output_without_heads)
 
-    def attention_scores(self, q, k, v):
+    def attention_logits(self, q, k, v):
 
         q, k, v = self.project_qkv(q, k, v)
 
@@ -301,8 +301,6 @@ class MHAttention(nn.Module):
         if self.causal:
             qk_scores.masked_fill_(self.mask, float("-inf"))
 
-        qk_scores = F.softmax(qk_scores, dim=-1)
-
         return qk_scores  # (batch, head, seq_len, seq_len)
 
     def reset_parameters(self):
@@ -529,15 +527,15 @@ class TransformerBlock(nn.Module):
 
         return x
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
         """
         Give back the attention scores used in this layer.
         """
         if self.pre_norm:
             x = self.layer_norm_1(x)
-            return self.attn.attention_scores(x, x, x)
+            return self.attn.attention_logits(x, x, x)
         else:
-            return self.attn.attention_scores(x, x, x)
+            return self.attn.attention_logits(x, x, x)
 
     def reset_parameters(self):
         self.layer_norm_1.reset_parameters()
@@ -697,7 +695,7 @@ class TransformerEncoder(nn.Module):
         else:
             return x
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
 
         x = self.preprocess(x)
 
@@ -705,8 +703,8 @@ class TransformerEncoder(nn.Module):
 
         for block in self.blocks:
             # Get attention scores with shape (batch, 1, head, seq_len, seq_len)
-            layer_attention_scores = block.attention_scores(x).unsqueeze(1)
-            layer_scores.append(layer_attention_scores)
+            layer_attention_logits = block.attention_logits(x).unsqueeze(1)
+            layer_scores.append(layer_attention_logits)
             x = block(x)
 
         return torch.cat(layer_scores, dim=1)  # (batch, layer, head, seq_len, seq_len)

broccoli/vit.py

@@ -400,9 +400,9 @@ class ViTEncoder(nn.Module):
     def forward(self, x):
         return self.encoder(x)
 
-    def attention_scores(self, x):
+    def attention_logits(self, x):
         x = self.encoder[:-1](x)
-        return self.encoder[-1].attention_scores(x)
+        return self.encoder[-1].attention_logits(x)
 
     def reset_parameters(self):
         for module in self.encoder:
@@ -546,16 +546,16 @@ class ViT(nn.Module):
     def forward(self, x):
         return self.pool(self.encoder(x))
 
-    def attention_scores(self, x):
-        return self.encoder.attention_scores(x)
+    def attention_logits(self, x):
+        return self.encoder.attention_logits(x)
 
     def head_to_bos_token_attention(self, x):
-        all_attention = self.attention_scores(x)
+        all_attention = self.attention_logits(x)
         batch_averages = torch.mean(all_attention, dim=0, keepdim=False)
         sequence_averages = torch.mean(batch_averages, dim=-1, keepdim=False)
         n_bos_tokens = self.encoder.encoder[-1]._bos_tokens
         just_bos = sequence_averages[:, :, :n_bos_tokens]
-        return F.softmax(just_bos, dim=-1)  # (layer, head, bos_token)
+        return F.softmax(just_bos, dim=1)  # (layer, head, bos_token)
 
     def reset_parameters(self):
         self.encoder.reset_parameters()

{broccoli_ml-5.1.2.dist-info → broccoli_ml-5.1.4.dist-info}/METADATA

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

{broccoli_ml-5.1.2.dist-info → broccoli_ml-5.1.4.dist-info}/RECORD

@@ -4,10 +4,10 @@ broccoli/cnn.py,sha256=WjoPDSpe3ttwxCBNfCVRdaCHvbeZ7G-a5_i8fUsK_d8,4889
 broccoli/linear.py,sha256=Y7s-DzcwsOipRboNHc4HTScw4mJRalNoVFsNcxOB6a4,4872
 broccoli/rope.py,sha256=GRqApBNmYCFaDak0WL1xE_BC5CTTYKQU_PBdeTcQcjc,12557
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=Gn8fhwgSq-izRc2nrsZ4JC5NsTIzQWFrol8pJ2pfzL4,23104
+broccoli/transformer.py,sha256=hLk4p-aS2xJiKHellkFKow5bgkAinoMlS8xTT5aQnII,23054
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=R73GTxx41FjVwAu4KDlYBAhytFo_9xVbbRbtBjAIW0s,20429
-broccoli_ml-5.1.2.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-5.1.2.dist-info/METADATA,sha256=iYvlWeqDeDaGDO_up_xC0NBLo3AhJOi6C_5BqWEgNYI,1368
-broccoli_ml-5.1.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-5.1.2.dist-info/RECORD,,
+broccoli/vit.py,sha256=Dy1egtEPft0l3812BENooc7SNogEibRPamXzy4jaJxA,20428
+broccoli_ml-5.1.4.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-5.1.4.dist-info/METADATA,sha256=s59fAtdmIjKhdP29kTPEgMp5buRGJzj5WgaBxsW83yg,1368
+broccoli_ml-5.1.4.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-5.1.4.dist-info/RECORD,,