broccoli-ml 3.1.2__tar.gz → 3.2.0__tar.gz

{broccoli_ml-3.1.2 → broccoli_ml-3.2.0}/PKG-INFO

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

{broccoli_ml-3.1.2 → broccoli_ml-3.2.0}/broccoli/rope.py

@@ -41,12 +41,12 @@ except ImportError:
     # Fallback: For PyTorch 1.6 to 1.9
     from torch.cuda.amp import autocast
 
-    def autocast_factory(_, enabled=True, dtype=None):
+    def autocast_factory(_, enabled=True):
         """
         A wrapper that mimics the modern autocast signature but calls the older
         torch.cuda.amp.autocast, ignoring the device_type argument.
         """
-        return autocast(enabled=enabled, dtype=dtype)
+        return autocast(enabled=enabled)
 
 
 from einops import rearrange, repeat

{broccoli_ml-3.1.2 → broccoli_ml-3.2.0}/broccoli/vit.py

@@ -58,6 +58,7 @@ class ClassificationHead(nn.Module):
         self,
         d_model,
         n_classes,
+        linear_module=nn.Linear,
         logit_projection_layer=nn.Linear,
         batch_norm_logits=True,
     ):
@@ -98,17 +99,19 @@ class SequencePoolClassificationHead(ClassificationHead):
         self,
         d_model,
         n_classes,
+        linear_module=nn.Linear,
         logit_projection_layer=nn.Linear,
         batch_norm_logits=True,
     ):
         super().__init__(
             d_model,
             n_classes,
+            linear_module=linear_module,
             logit_projection_layer=logit_projection_layer,
             batch_norm_logits=batch_norm_logits,
         )
 
-        self.summarize = SequencePool(d_model, logit_projection_layer)
+        self.summarize = SequencePool(d_model, linear_module)
         # Rebuild the classification process with the correct summary module:
         self.classification_process = nn.Sequential(
             *[
@@ -510,6 +513,7 @@ class ViT(nn.Module):
         self.pool = head(
             transformer_embedding_size,
             image_classes,
+            linear_module=linear_module,
             logit_projection_layer=logit_projection_layer,
             batch_norm_logits=batch_norm_logits,
         )

{broccoli_ml-3.1.2 → broccoli_ml-3.2.0}/pyproject.toml

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