braindecode 1.5.0.dev984__py3-none-any.whl → 1.5.0.dev182195895__py3-none-any.whl

braindecode/models/biot.py

@@ -439,7 +439,7 @@ class _BIOTEncoder(nn.Module):
         self.channel_tokens = nn.Embedding(
             num_embeddings=n_chans, embedding_dim=emb_size
         )
-        self.index = nn.Parameter(torch.LongTensor(range(n_chans)), requires_grad=False)
+        self.register_buffer("index", torch.arange(n_chans, dtype=torch.long))
 
     def stft(self, sample):
         """

braindecode/modules/attention.py

@@ -837,6 +837,18 @@ class CATLite(nn.Module):
 class MultiHeadAttention(nn.Module):
     """Multi-head self-attention block.
 
+    Uses ``F.scaled_dot_product_attention`` for optimized attention
+    kernels (flash-attention on CUDA, memory-efficient on other devices).
+
+    Parameters
+    ----------
+    emb_size : int
+        The embedding dimension.
+    num_heads : int
+        Number of attention heads. Must evenly divide ``emb_size``.
+    dropout : float, optional
+        Dropout probability applied to attention weights. Default: 0.0.
+
     Examples
     --------
     >>> import torch
@@ -848,40 +860,57 @@ class MultiHeadAttention(nn.Module):
     torch.Size([2, 10, 32])
     """
 
-    def __init__(self, emb_size, num_heads, dropout):
+    def __init__(self, emb_size, num_heads, dropout=0.0):
         super().__init__()
+        if emb_size % num_heads != 0:
+            raise ValueError(
+                f"emb_size ({emb_size}) must be divisible by num_heads ({num_heads})."
+            )
         self.emb_size = emb_size
         self.num_heads = num_heads
+        self.head_dim = emb_size // num_heads
         self.keys = nn.Linear(emb_size, emb_size)
         self.queries = nn.Linear(emb_size, emb_size)
         self.values = nn.Linear(emb_size, emb_size)
-        self.att_drop = nn.Dropout(dropout)
+        self.att_drop = dropout
         self.projection = nn.Linear(emb_size, emb_size)
 
         self.rearrange_stack = Rearrange(
-            "b n (h d) -> b h n d",
-            h=num_heads,
+            "batch seq (heads head_dim) -> batch heads seq head_dim",
+            heads=num_heads,
         )
         self.rearrange_unstack = Rearrange(
-            "b h n d -> b n (h d)",
+            "batch heads seq head_dim -> batch seq (heads head_dim)",
         )
 
     def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor:
+        """Forward pass.
+
+        Parameters
+        ----------
+        x : Tensor
+            Input tensor of shape ``(batch, seq, emb_size)``.
+        mask : Tensor, optional
+            Attention mask following PyTorch SDPA convention: for boolean
+            masks ``True`` means **ignore** that position; for float
+            masks the values are **added** to attention scores before
+            softmax.
+        """
         queries = self.rearrange_stack(self.queries(x))
         keys = self.rearrange_stack(self.keys(x))
         values = self.rearrange_stack(self.values(x))
-        energy = torch.einsum("bhqd, bhkd -> bhqk", queries, keys)
-        if mask is not None:
-            fill_value = float("-inf")
-            energy = energy.masked_fill(~mask, fill_value)
-
-        scaling = self.emb_size ** (1 / 2)
-        att = F.softmax(energy / scaling, dim=-1)
-        att = self.att_drop(att)
-        out = torch.einsum("bhal, bhlv -> bhav ", att, values)
+
+        dp = self.att_drop if self.training else 0.0
+        out = F.scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            attn_mask=mask,
+            dropout_p=dp,
+        )
+
         out = self.rearrange_unstack(out)
-        out = self.projection(out)
-        return out
+        return self.projection(out)
 
 
 class CrissCrossTransformerEncoderLayer(nn.Module):

braindecode/version.py

@@ -1 +1 @@
-__version__ = "1.5.0.dev984"
+__version__ = "1.5.0.dev182195895"

{braindecode-1.5.0.dev984.dist-info → braindecode-1.5.0.dev182195895.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.5.0.dev984
+Version: 1.5.0.dev182195895
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>

{braindecode-1.5.0.dev984.dist-info → braindecode-1.5.0.dev182195895.dist-info}/RECORD

@@ -3,7 +3,7 @@ braindecode/classifier.py,sha256=7kC_oY_UzHEes_WWdCvEpiA1ZKxMeuLL5tIPp5rfcpg,962
 braindecode/eegneuralnet.py,sha256=xjE6aPZdCQPs29NIpy_m1GLMMC2WZ3Db0Fuh1-xE1h4,13827
 braindecode/regressor.py,sha256=KiMJpqCUPWA2k2JWk9HGYTzeoBqJ4gAKEudeUVcFZY4,9266
 braindecode/util.py,sha256=f8bNIwt-SwsHqheH_BADQxTtA9oPt3Lb7GFnoI-Huwc,14101
-braindecode/version.py,sha256=5WqXRAeL_e0Aq78aGQsveuaWwGcV_OgtE7s9XQXznkc,29
+braindecode/version.py,sha256=FCjD8jntd2wdAgufj5JsChb_O65vD4HdEFpKyVmbAz0,35
 braindecode/augmentation/__init__.py,sha256=4xune2QUK6KHMKsAqijF7I9eeiVbP0wEoQJjCNLNcKM,1081
 braindecode/augmentation/base.py,sha256=OJ1shOljI1yTY9zh2qWxQwivlY43sfx9Q-MAyMhxtPs,7338
 braindecode/augmentation/functional.py,sha256=q2k6mAXrujYlOZUndcjZN8e8b-6oJF1gGsORAI23hyE,43998
@@ -45,7 +45,7 @@ braindecode/models/attentionbasenet.py,sha256=k7ar7aEjANudPu7krAZsRx-ag61ugirS1X
 braindecode/models/attn_sleep.py,sha256=F9x4spTtzfiCC1h9UYITmIDQeJW6_2CXTZktZX9R0RE,17950
 braindecode/models/base.py,sha256=yGJgr0f5rD-gJZ5Msw9FzGzVO-x_re-6jOSK8Iht6x4,27923
 braindecode/models/bendr.py,sha256=kqIKtIrgHhIdwYfBFwAJ1g69YbdpA8bs2C2e2GUT0rY,25643
-braindecode/models/biot.py,sha256=k742atP5TMAFkEDQ1hd-AhOJW2a0R6_StRbZGYjWS1U,19358
+braindecode/models/biot.py,sha256=EJuf77ieWirrkYrVZessQF-MulVBdNTUhpRw3UxHhb8,19348
 braindecode/models/brainmodule.py,sha256=idyQVTp3VBJXKF1YjMx8o1kUKrcL_E_AJsioPjcuqV8,33282
 braindecode/models/cbramod.py,sha256=ZNyM_iBGFGjMpmOiqhLwmgGzy-AowBJqTORzNeNG99I,14309
 braindecode/models/config.py,sha256=9yc7fh2p6txrjaOROyJevJO7ud5fofZk6C76JfwZrX0,8431
@@ -93,7 +93,7 @@ braindecode/models/usleep.py,sha256=VaIRbTSuKisD_W6fMwLlY5bNiH41fBnfx0k-iiLHWpw,
 braindecode/models/util.py,sha256=QSku9ZlXTyJzMuTLG-VlPPifTe2_6huo7i9-zQnhZho,17728
 braindecode/modules/__init__.py,sha256=NChqITh6yMcagoYFVkQC14aJjCAGkRO8fYQrH_0Zv1k,1977
 braindecode/modules/activation.py,sha256=QSlpMj_mfroLvaaIGfnjWJbWgvpEBigq2AU0jj0fglQ,2496
-braindecode/modules/attention.py,sha256=Rvad4FchhXSjgApi15aQTXRsWYJ43aChOH-akSbRJAI,32689
+braindecode/modules/attention.py,sha256=zHJ6kL54ZGGX4376Xk2rzxRMateOi4AtSSxzLs5hD-o,33629
 braindecode/modules/blocks.py,sha256=Nv-hu1nx0zZBs0aYMe7FDCzRkH1rf34CHXpoTQP6lCE,4993
 braindecode/modules/convolution.py,sha256=b6jvSp1bS8skIUfX2Ztn_JHyxX8q3hvaEqyOmzNnRsc,10786
 braindecode/modules/filter.py,sha256=Lcq5zt3uEmIiFeu8SbAWkhsqIoSU1zS_Mo-0mBFHrJE,25918
@@ -119,9 +119,9 @@ braindecode/training/scoring.py,sha256=WRkwqbitA3m_dzRnGp2ZIZPge5Nhx9gAEQhIHzeH4
 braindecode/visualization/__init__.py,sha256=4EER_xHqZIDzEvmgUEm7K1bgNKpyZAIClR9ZCkMuY4M,240
 braindecode/visualization/confusion_matrices.py,sha256=qIWMLEHow5CJ7PhGggD8mnD55Le6xhma9HSzt4R33fc,9509
 braindecode/visualization/gradients.py,sha256=-NuWIlpOeSR2OWGfUl6jn1RqdNPxmiQNLPCckWb8SRE,2330
-braindecode-1.5.0.dev984.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
-braindecode-1.5.0.dev984.dist-info/licenses/NOTICE.txt,sha256=sOxuTbalPxTM8H6VqtvGbXCt_BoOF7JevEYG_knqbm4,620
-braindecode-1.5.0.dev984.dist-info/METADATA,sha256=gs4siQJaVjCx8UszxeuPfWYFF_FPk2aagFqBzmIZbTI,10099
-braindecode-1.5.0.dev984.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-braindecode-1.5.0.dev984.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
-braindecode-1.5.0.dev984.dist-info/RECORD,,
+braindecode-1.5.0.dev182195895.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
+braindecode-1.5.0.dev182195895.dist-info/licenses/NOTICE.txt,sha256=sOxuTbalPxTM8H6VqtvGbXCt_BoOF7JevEYG_knqbm4,620
+braindecode-1.5.0.dev182195895.dist-info/METADATA,sha256=VSuDNooSSWCZExCgShqwhOq2EayUOJQ5w_bixsGHC8A,10105
+braindecode-1.5.0.dev182195895.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+braindecode-1.5.0.dev182195895.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
+braindecode-1.5.0.dev182195895.dist-info/RECORD,,