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

broccoli/transformer.py

@@ -535,9 +535,9 @@ class TransformerBlock(nn.Module):
         """
         if self.pre_norm:
             x = self.layer_norm_1(x)
-            return self.attn(x, x, x)
+            return self.attn.attention_scores(x, x, x)
         else:
-            return self.attn(x, x, x)
+            return self.attn.attention_scores(x, x, x)
 
     def reset_parameters(self):
         self.layer_norm_1.reset_parameters()
@@ -683,6 +683,8 @@ class TransformerEncoder(nn.Module):
                 )  # to shape (1, seq_len) to broadcast over batch
             )
 
+        return x
+
     def forward(self, x):
 
         x = self.preprocess(x)

broccoli/vit.py

@@ -553,7 +553,7 @@ class ViT(nn.Module):
         all_attention = self.attention_scores(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._bos_tokens
+        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)
 

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

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

{broccoli_ml-5.1.0.dist-info → broccoli_ml-5.1.2.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=x3Mo6_1x6fGG6lPDPx9srxn6UdwKEpvjFAO8zoMwAMI,23052
+broccoli/transformer.py,sha256=Gn8fhwgSq-izRc2nrsZ4JC5NsTIzQWFrol8pJ2pfzL4,23104
 broccoli/utils.py,sha256=oOWzn6dJ5nC_9r4zq0emmfmaYACJXJNFS48AOpW2jqc,358
-broccoli/vit.py,sha256=tUYQyoDsBc5ZR_M5_J0huj0T3OAy-vn1f19hCGVDCrM,20425
-broccoli_ml-5.1.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-5.1.0.dist-info/METADATA,sha256=3986lqn1iuWJ53O8ckM9LVU3tTjr32i19SeIXauWDXw,1368
-broccoli_ml-5.1.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-5.1.0.dist-info/RECORD,,
+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,,