broccoli-ml 0.36.0__py3-none-any.whl → 0.38.0__py3-none-any.whl

broccoli/linear.py

@@ -131,8 +131,7 @@ class WeightNormedLinear(nn.Module):
         return F.linear(x, self.weights(), self.bias)
 
     def __repr__(self) -> str:
-        # Optional: A nice representation for printing the module.
         return (
-            f"AnchoredLinear(in_features={self.in_features},"
+            f"WeightNormedLinear(in_features={self.in_features},"
             f"out_features={self.out_features}, bias={self.use_bias})"
         )

broccoli/transformer.py

@@ -13,6 +13,45 @@ from .rope import RotaryEmbedding, apply_rotary_emb
 from .linear import AnchoredLinear, SpectralNormLinear
 
 
+def drop_path(
+    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
+):
+    """
+    From https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py
+    Copyright 2019 Ross Wightman
+    See documentation and licence there.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """
+    From https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py
+    Copyright 2019 Ross Wightman
+    See documentation and licence there.
+    """
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+    def extra_repr(self):
+        return f"drop_prob={round(self.drop_prob, 3):0.3f}"
+
+
 class MHAttention(nn.Module):
     """
     Multi-head self-attention using einops and optionally a custom linear layer.
@@ -264,6 +303,8 @@ class TransformerBlock(nn.Module):
         mlp_ratio=4,
         activation: nn.Module = nn.ReLU,
         activation_kwargs: Optional[dict] = None,
+        ff_linear_module_up=None,
+        ff_linear_module_down=None,
         mlp_dropout=0.0,
         msa_dropout=0.0,
         identity_probability=0.0,
@@ -279,10 +320,11 @@ class TransformerBlock(nn.Module):
         self.post_norm = post_norm
         self.normformer = normformer
 
-        self.identity_probability = identity_probability
+        self.drop_path = DropPath(drop_prob=identity_probability, scale_by_keep=True)
 
         self.layer_norm_1 = nn.LayerNorm(d_model)
         self.layer_norm_2 = nn.LayerNorm(d_model)
+        self.layer_norm_3 = nn.LayerNorm(d_model)
 
         if position_embedding_type == "relative":
             max_freq = int(max(source_size) / 2)  # Suggested by Gemini!
@@ -316,12 +358,20 @@ class TransformerBlock(nn.Module):
             activation=activation,
             activation_kwargs=activation_kwargs,
             dropout=mlp_dropout,
-            linear_module_up=linear_module,
-            linear_module_down=linear_module,
-            pre_norm=pre_norm,
+            linear_module_up=(
+                ff_linear_module_up
+                if ff_linear_module_up is not None
+                else linear_module
+            ),
+            linear_module_down=(
+                ff_linear_module_down
+                if ff_linear_module_down is not None
+                else linear_module
+            ),
+            pre_norm=False,  # Handled outside the block
             normformer=normformer,
-            post_norm=post_norm,
-            residual_path=True,
+            post_norm=False,  # Handled outside the block
+            residual_path=False,  # Handled outside the block
         )
 
     @property
@@ -329,34 +379,23 @@ class TransformerBlock(nn.Module):
         return self.attn._kv_distance
 
     def forward(self, x):
-        if not self.training:
-            identity_probability = 0.0
-        else:
-            identity_probability = self.identity_probability
-
-        # perform the identity operation for some rows in the batch
-        dist = torch.distributions.Binomial(x.size(0), identity_probability)
-        identity_count = int(dist.sample().item())
-
-        shuffle_indices = torch.randperm(x.size(0), device=x.device)
-        unshuffle_indices = torch.argsort(shuffle_indices)
-        shuffled = x[shuffle_indices, :, :]
-        identity_x = shuffled[:identity_count, :, :]
-        process_x = shuffled[identity_count:, :, :]
-
-        residual_x = process_x
 
         if self.pre_norm:
-            process_x = self.layer_norm_1(process_x)
-
-        process_x = residual_x + self.attn(process_x, process_x, process_x)
-
-        if self.post_norm:
-            process_x = self.layer_norm_2(process_x)
-
-        process_x = self.ff(process_x)
+            normx = self.layer_norm_1(x)
+            x = x + self.drop_path(self.attn(normx, normx, normx))
+            normx = self.layer_norm_2(x)
+            x = x + self.drop_path(self.ff(normx))
+        elif self.post_norm:
+            x = x + self.drop_path(self.attn(x, x, x))
+            x = self.layer_norm_1(x)
+            x = x + self.drop_path(self.ff(x))
+            x = self.layer_norm_2(x)
+        else:
+            x = x + self.drop_path(self.attn(x, x, x))
+            x = x + self.drop_path(self.ff(x))
 
-        x = torch.cat([identity_x, process_x])[unshuffle_indices, :, :].contiguous()
+        if self.pre_norm and self.post_norm:
+            x = self.layer_norm_3(x)
 
         return x
 
@@ -378,6 +417,8 @@ class TransformerEncoder(nn.Module):
         mlp_ratio=4,
         activation: nn.Module = nn.ReLU,
         activation_kwargs: Optional[dict] = None,
+        ff_linear_module_up=None,
+        ff_linear_module_down=None,
         mlp_dropout=0.0,
         msa_dropout=0.0,
         stochastic_depth=0.0,
@@ -441,6 +482,8 @@ class TransformerEncoder(nn.Module):
                     mlp_ratio=mlp_ratio,
                     activation=activation,
                     activation_kwargs=activation_kwargs,
+                    ff_linear_module_up=ff_linear_module_up,
+                    ff_linear_module_down=ff_linear_module_down,
                     mlp_dropout=mlp_dropout,
                     msa_dropout=msa_dropout,
                     identity_probability=self.stochastic_depth_probabilities[i],

broccoli/vit.py

@@ -132,6 +132,8 @@ class ViTEncoder(nn.Module):
         transformer_return_bos_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
+        transformer_ff_linear_module_up=None,
+        transformer_ff_linear_module_down=None,
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
@@ -282,6 +284,8 @@ class ViTEncoder(nn.Module):
                 mlp_ratio=transformer_mlp_ratio,
                 activation=transformer_activation,
                 activation_kwargs=transformer_activation_kwargs,
+                ff_linear_module_up=transformer_ff_linear_module_up,
+                ff_linear_module_down=transformer_ff_linear_module_down,
                 mlp_dropout=transformer_mlp_dropout,
                 msa_dropout=transformer_msa_dropout,
                 stochastic_depth=transformer_stochastic_depth,
@@ -305,14 +309,16 @@ class ViTEncoder(nn.Module):
                 activation_kwargs=transformer_activation_kwargs,
                 dropout=transformer_mlp_dropout,
                 linear_module_up=(
+                    # First truthy assigned value
                     transformer_initial_ff_linear_module_up
-                    if transformer_initial_ff_linear_module_up is not None
-                    else linear_module
+                    or transformer_ff_linear_module_up
+                    or linear_module
                 ),
                 linear_module_down=(
+                    # First truthy assigned value
                     transformer_initial_ff_linear_module_down
-                    if transformer_initial_ff_linear_module_down is not None
-                    else linear_module
+                    or transformer_ff_linear_module_down
+                    or linear_module
                 ),
                 pre_norm=transformer_pre_norm,
                 normformer=transformer_normformer,
@@ -386,6 +392,8 @@ class ViT(nn.Module):
         transformer_return_bos_tokens=False,
         transformer_activation: nn.Module = SquaredReLU,
         transformer_activation_kwargs: Optional[dict] = None,
+        transformer_ff_linear_module_up=None,
+        transformer_ff_linear_module_down=None,
         transformer_mlp_dropout=0.0,
         transformer_msa_dropout=0.1,
         transformer_stochastic_depth=0.1,
@@ -446,6 +454,8 @@ class ViT(nn.Module):
             transformer_return_bos_tokens=transformer_return_bos_tokens,
             transformer_activation=transformer_activation,
             transformer_activation_kwargs=transformer_activation_kwargs,
+            transformer_ff_linear_module_up=transformer_ff_linear_module_up,
+            transformer_ff_linear_module_down=transformer_ff_linear_module_down,
             transformer_mlp_dropout=transformer_mlp_dropout,
             transformer_msa_dropout=transformer_msa_dropout,
             transformer_stochastic_depth=transformer_stochastic_depth,

{broccoli_ml-0.36.0.dist-info → broccoli_ml-0.38.0.dist-info}/METADATA

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

{broccoli_ml-0.36.0.dist-info → broccoli_ml-0.38.0.dist-info}/RECORD

@@ -5,13 +5,13 @@ broccoli/assets/cifar100_eigenvectors_size_2.pt,sha256=DjXDOXMeuMpIqNuGhX9z-OWYV
 broccoli/assets/cifar100_eigenvectors_size_3.pt,sha256=gL6k0xtXYiYP6ZSvEiMBdJ7kIkT0AngTpDJHFQqwgxA,7173
 broccoli/cnn.py,sha256=jeRyKIAMWu1E3iyI14MGgSZuZivPMh12iqkqW9ilNjo,17785
 broccoli/eigenpatches.py,sha256=J6n2usN1oQuHEHYiBNyYpn_a9eQcHjOBiIlvSei520Y,2413
-broccoli/linear.py,sha256=8Y9vD85ZEgNZsIQgO3uRQ3lOQR-JjwvabY8liCrfNCk,4831
+broccoli/linear.py,sha256=zlFDij9TngqDzTTpUlZtX0PXAQgxyWWnTIXpiO1rBk0,4768
 broccoli/rope.py,sha256=hw7kBPNR9GQXj4GxyIAffsGKPfcTPOFh8Bc7oEHtaZY,12108
 broccoli/tensor.py,sha256=um8mrxkYbvNDo-QvHlmJm8Aw6qcngOlUZPoAk_PMReA,4480
-broccoli/transformer.py,sha256=NH94U6lxHzmDGDHTTtJV2kUs7IcS2iNmFJl44_6KtQ0,15456
+broccoli/transformer.py,sha256=Xw1oBLsvVeHmMqgurhorRa49nrjfooLev5uBPFeK9og,17004
 broccoli/utils.py,sha256=htq_hOsdhUhL0nJi9WkKiEYOjEoWqFpK5X49PtgTf-0,299
-broccoli/vit.py,sha256=05xqIw9xvE5easXcp4wIA1jQ0xUyRIq6h0ZDtbitXi4,17184
-broccoli_ml-0.36.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
-broccoli_ml-0.36.0.dist-info/METADATA,sha256=csog4ZG1PGeRuFO5QnHdVPgmDYXsGQQJ621JgU0D83w,1257
-broccoli_ml-0.36.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-broccoli_ml-0.36.0.dist-info/RECORD,,
+broccoli/vit.py,sha256=RaSJZh2ogqNAvAkQDuZpNqdtRTWaW_8ug4BsBCBK_f4,17728
+broccoli_ml-0.38.0.dist-info/LICENSE,sha256=0BAzJE5BqQ7Iixp_AFdB2W1uO-HCRX-Qfun8PHt6yVM,1073
+broccoli_ml-0.38.0.dist-info/METADATA,sha256=8lcwJvmlPT_-SkJZ5Qd6nFc4Njz7mQD-JefeUxhk8Xw,1257
+broccoli_ml-0.38.0.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+broccoli_ml-0.38.0.dist-info/RECORD,,