Braindecode 1.8.1.dev168512945__tar.gz → 1.8.1.dev168570878__tar.gz

{braindecode-1.8.1.dev168512945/braindecode.egg-info → braindecode-1.8.1.dev168570878}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev168512945
+Version: 1.8.1.dev168570878
 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.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/functional/__init__.py

@@ -1,4 +1,3 @@
-from .attention import scaled_dot_product_attention
 from .functions import (
     _get_gaussian_kernel1d,
     apply_rotary,
@@ -46,7 +45,6 @@ __all__ = [
     "rotate_half",
     "rotate_pairs",
     "safe_log",
-    "scaled_dot_product_attention",
     "sinusoidal_positional_encoding",
     "spectral_input",
     "square",

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/functional/functions.py

@@ -588,7 +588,6 @@ def sinusoidal_positional_encoding(
     dim: int,
     device: torch.device | None = None,
     dtype: torch.dtype = torch.float32,
-    base: float = 10000.0,
 ) -> torch.Tensor:
     r"""Fixed sine/cosine positional-encoding table of shape ``(n_positions, dim)``.
 
@@ -610,9 +609,6 @@ def sinusoidal_positional_encoding(
         Device of the table (default: CPU), for tables built inside ``forward``.
     dtype : torch.dtype
         Floating dtype the table is computed in.
-    base : float, default=10000.0
-        Longest wavelength, the ``10000`` of Vaswani et al.; SignalJEPA uses the
-        largest number of time samples it encodes.
 
     Returns
     -------
@@ -624,7 +620,7 @@ def sinusoidal_positional_encoding(
     position = torch.arange(n_positions, device=device).unsqueeze(1).to(dtype)
     div_term = torch.exp(
         torch.arange(0, dim_even, 2, device=device).to(dtype)
-        * (-math.log(base) / dim_even)
+        * (-math.log(10000.0) / dim_even)
     )
     pe = torch.zeros(n_positions, dim_even, device=device, dtype=dtype)
     pe[:, 0::2] = torch.sin(position * div_term)

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/atcnet.py

@@ -1,6 +1,7 @@
 # Authors: Cedric Rommel <cedric.rommel@inria.fr>
 #
 # License: BSD (3-clause)
+import math
 from typing import Optional
 
 import torch
@@ -9,7 +10,6 @@ from mne.utils import warn
 from torch import nn
 from torch.nn.utils.parametrize import is_parametrized, register_parametrization
 
-from braindecode.functional import scaled_dot_product_attention
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import (
     CausalConv1d,
@@ -934,10 +934,20 @@ class _MHA(nn.Module):
         K_ = torch.cat(K.split(self.head_dim, -1), 0)  # (B', S, D')
         V_ = torch.cat(V.split(self.head_dim, -1), 0)  # (B', S, D')
 
-        # Multihead output (batch_size, seq_len, dim): softmax(Q K^T / sqrt(d)) V
-        # per head, heads merged back along the feature axis.
+        # Attention weights of size (num_heads * batch_size, n, m):
+        # measures how similar each pair of Q and K is.
+        W = torch.softmax(
+            Q_.bmm(K_.transpose(-2, -1)) / math.sqrt(self.head_dim),
+            -1,  # (B', D', S)
+        )  # (B', N, M)
+
+        # Multihead output (batch_size, seq_len, dim):
+        # weighted sum of V where a value gets more weight if its corresponding
+        # key has larger dot product with the query.
         H = torch.cat(
-            scaled_dot_product_attention(Q_, K_, V_, exact=True).split(batch_size, 0),
+            (W.bmm(V_)).split(  # (B', S, S)  # (B', S, D')
+                batch_size, 0
+            ),  # [(B, S, D')] * num_heads
             -1,
         )  # (B, S, D)
 

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/attn_sleep.py

@@ -27,14 +27,15 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import math
 import warnings
 from copy import deepcopy
 from numbers import Integral
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
-from braindecode.functional import scaled_dot_product_attention
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import CausalConv1d
 
@@ -410,6 +411,14 @@ class _MRCNN(nn.Module):
 ##########################################################################################
 
 
+def _attention(query: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
+    """Scaled dot product attention weights."""
+    # d_k - dimension of the query and key vectors
+    d_k = query.size(-1)
+    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
+    return F.softmax(scores, dim=-1)  # attention weights
+
+
 class _MultiHeadedAttention(nn.Module):
     def __init__(self, h, d_model, after_reduced_cnn_size, dropout=0.1):
         """Take in model size and number of heads."""
@@ -455,14 +464,9 @@ class _MultiHeadedAttention(nn.Module):
             .transpose(1, 2)
         )
 
-        # exact: the reference op order, same outputs as its hand-written attention
-        x = scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            dropout_p=self.dropout.p if self.training else 0.0,
-            exact=True,
-        )
+        # apply dropout to the *weights*, then the weighted sum
+        attn = self.dropout(_attention(query, key))
+        x = torch.matmul(attn, value)
 
         # merge heads and project
         x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_per_head)

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/eegdino.py

@@ -10,7 +10,7 @@ from typing import Dict, Optional, Sequence, Union
 import torch
 from torch import nn
 
-from braindecode.functional import scaled_dot_product_attention, spectral_input
+from braindecode.functional import spectral_input
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import warn_if_sfreq_differs
 from braindecode.modules import DropPath, PatchTokenizer
@@ -500,16 +500,9 @@ class _Attention(nn.Module):
             .permute(2, 0, 3, 1, 4)
             .unbind(0)
         )
-        # exact with q pre-scaled: the reference op order (q * scale) @ k^T
-        out = scaled_dot_product_attention(
-            query * self.scale,
-            key,
-            value,
-            dropout_p=self.attn_drop.p if self.training else 0.0,
-            scale=1.0,
-            exact=True,
-        )
-        out = out.transpose(1, 2).flatten(2)
+        attn = (query * self.scale @ key.transpose(-2, -1)).softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        out = (attn @ value).transpose(1, 2).flatten(2)
         return self.proj_drop(self.proj(out))
 
 

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/popt.py

@@ -12,6 +12,7 @@
 
 from __future__ import annotations
 
+import math
 import warnings
 from collections import OrderedDict
 
@@ -19,7 +20,6 @@ import numpy as np
 import torch
 import torch.nn as nn
 
-from braindecode.functional import sinusoidal_positional_encoding
 from braindecode.models.base import EEGModuleMixin
 from braindecode.models.util import extract_channel_locations_from_chs_info
 
@@ -355,7 +355,13 @@ class _PopTSpatialPositionalEncoding(nn.Module):
                 "axis and the sequence id)."
             )
         pe_dim = hidden_dim // 4
-        pe = sinusoidal_positional_encoding(max_len, pe_dim)
+        pe = torch.zeros(max_len, pe_dim)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, pe_dim, 2).float() * (-math.log(10000.0) / pe_dim)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
         self.register_buffer("pe", pe.unsqueeze(0), persistent=False)
 
     def forward(

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/signal_jepa.py

@@ -13,7 +13,7 @@ import torch
 from einops.layers.torch import Rearrange
 from torch import nn
 
-from braindecode.functional import init_weights_, sinusoidal_positional_encoding
+from braindecode.functional import init_weights_
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules.attention import (
     _hpu_lazy,
@@ -1682,9 +1682,6 @@ def _pos_encode_time(
 ):
     """1-dimensional positional encoding.
 
-    :func:`~braindecode.functional.sinusoidal_positional_encoding` with the
-    largest number of time samples as the longest wavelength.
-
     Parameters
     ----------
         n_times: int
@@ -1701,9 +1698,14 @@ def _pos_encode_time(
         pos_encoding: (n_times, n_dim)
     """
     assert n_dim % 2 == 0
-    return sinusoidal_positional_encoding(
-        n_times, n_dim, device=device, base=float(max_n_times)
+    position = torch.arange(n_times, device=device).unsqueeze(1)
+    div_term = torch.exp(
+        torch.arange(0, n_dim, 2, device=device) * (-math.log(max_n_times) / n_dim)
     )
+    pos_encoding = torch.empty((n_times, n_dim), dtype=torch.float32, device=device)
+    pos_encoding[:, 0::2] = torch.sin(position * div_term)
+    pos_encoding[:, 1::2] = torch.cos(position * div_term)
+    return pos_encoding
 
 
 def _pos_encode_contineous(

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/tcformer.py

@@ -16,7 +16,6 @@ from braindecode.functional import (
     init_weights_,
     rotary_angles,
     rotary_frequencies,
-    scaled_dot_product_attention,
 )
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import CausalConv1d, Conv1dWithConstraint, DropPath
@@ -368,16 +367,10 @@ class _GroupedQueryAttention(nn.Module):
         value = value.repeat_interleave(groups_per_kv, dim=1)
         query = apply_rotary(query, cos[:seq_len], sin[:seq_len])
         key = apply_rotary(key, cos[:seq_len], sin[:seq_len])
-        # exact: the reference op order (q @ k^T) * scale
-        context = scaled_dot_product_attention(
-            query,
-            key,
-            value,
-            dropout_p=self.drop.p if self.training else 0.0,
-            scale=self.scale,
-            exact=True,
-        )
-        context = context.transpose(1, 2).flatten(2)
+        key_t = key.transpose(-2, -1)
+        attention = ((query @ key_t) * self.scale).softmax(dim=-1)
+        attention = self.drop(attention)
+        context = (attention @ value).transpose(1, 2).flatten(2)
         return self.out_proj(context)
 
 

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/models/tmsanet.py

@@ -8,11 +8,13 @@
 
 from __future__ import annotations
 
+import math
+
 import torch
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
+from torch.nn import functional as F
 
-from braindecode.functional import scaled_dot_product_attention
 from braindecode.models.base import EEGModuleMixin
 from braindecode.modules import FeedForwardBlock
 
@@ -203,14 +205,12 @@ class _TMSAAttention(nn.Module):
         k_local = self.split_heads(self.w_k_local(local))
         k_global = self.split_heads(self.w_k_global(x))
         v = self.split_heads(self.w_v(x))
-        # exact: the reference op order and its dropout RNG order
-        dropout_p = self.att_drop.p if self.training else 0.0
-        x_local = scaled_dot_product_attention(
-            q, k_local, v, dropout_p=dropout_p, exact=True
-        )
-        x_global = scaled_dot_product_attention(
-            q, k_global, v, dropout_p=dropout_p, exact=True
-        )
+        # Explicit softmax attention keeps the reference's dropout RNG order.
+        scale = math.sqrt(self.head_dim)
+        attn_local = F.softmax(q @ k_local.transpose(-2, -1) / scale, dim=-1)
+        x_local = self.att_drop(attn_local) @ v
+        attn_global = F.softmax(q @ k_global.transpose(-2, -1) / scale, dim=-1)
+        x_global = self.att_drop(attn_global) @ v
         return self.w_o(self.merge_heads(x_local + x_global))
 
 

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode/modules/attention.py

@@ -16,19 +16,12 @@ import math
 from typing import Optional
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange
 from einops.layers.torch import Rearrange
 from torch import Tensor, nn
 
 from braindecode.functional import _get_gaussian_kernel1d
-from braindecode.functional.attention import (  # noqa: F401  (re-exported)
-    _hpu_eager_fused_sdpa,
-    _hpu_lazy,
-    _sdpa_math_dropout,
-)
-from braindecode.functional.attention import (
-    scaled_dot_product_attention as _scaled_dot_product_attention,
-)
 
 
 class SqueezeAndExcitation(nn.Module):
@@ -841,6 +834,91 @@ class CATLite(nn.Module):
         return channel_score * x
 
 
+@torch.jit.unused
+def _hpu_eager_fused_sdpa(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    scale: Optional[float] = None,
+    fused_in_lazy: bool = False,
+) -> Optional[Tensor]:
+    # Lazy mode: None, the graph compiler already fuses torch's SDPA there
+    # (unless the caller measured the fused kernel faster for its shapes).
+    from habana_frameworks.torch.utils.internal import is_lazy
+
+    if is_lazy() and not fused_in_lazy:
+        return None
+    from habana_frameworks.torch.hpex.kernels import FusedSDPA
+
+    # Strided q/k/v views (heads split off a projection) give wrong outputs and
+    # gradients at some shapes; contiguous copies give the right result.
+    query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
+    return FusedSDPA.apply(query, key, value, None, 0.0, False, scale, "None", False)
+
+
+@torch.jit.unused
+def _hpu_lazy() -> bool:
+    from habana_frameworks.torch.utils.internal import is_lazy
+
+    return is_lazy()
+
+
+def _sdpa_math_dropout(
+    query: Tensor, key: Tensor, value: Tensor, dropout_p: float
+) -> Tensor:
+    # torch's math path for SDPA with dropout, op for op, without the pass of
+    # ``_safe_softmax`` that zeroes fully masked rows (none without a mask):
+    # same values, same dropout draws.
+    scale = math.sqrt(1.0 / math.sqrt(query.size(-1)))
+    attn = torch.matmul(query * scale, key.transpose(-2, -1) * scale)
+    attn = F.dropout(attn.softmax(-1), dropout_p, True)
+    return torch.matmul(attn, value)
+
+
+def _scaled_dot_product_attention(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_mask: Optional[Tensor] = None,
+    dropout_p: float = 0.0,
+    fused_in_lazy: bool = False,
+) -> Tensor:
+    """``F.scaled_dot_product_attention``; Gaudi's fused kernel on HPU.
+
+    In HPU eager mode, torch's SDPA runs op by op (the full
+    ``(batch, heads, seq, seq)`` scores in device memory), 2-3.5x slower than
+    the bridge's fused kernel on long sequences. The fused kernel is used
+    without a mask or dropout (its dropout would draw other masks). In lazy
+    mode only with ``fused_in_lazy=True``: the graph compiler fuses torch's
+    SDPA too, faster on long sequences, slower on short ones.
+    """
+    if (
+        not torch.jit.is_scripting()
+        and query.device.type == "hpu"
+        and attn_mask is None
+        and dropout_p == 0.0
+        and (query.dtype == torch.float32 or query.dtype == torch.bfloat16)
+    ):
+        out = _hpu_eager_fused_sdpa(query, key, value, fused_in_lazy=fused_in_lazy)
+        if out is not None:
+            return out
+    if (
+        not torch.jit.is_scripting()
+        and dropout_p > 0.0
+        and attn_mask is None
+        and (query.dtype == torch.float32 or query.dtype == torch.float64)
+        and (
+            query.device.type == "cpu"
+            or (query.device.type == "hpu" and not _hpu_lazy())
+        )
+    ):
+        # CPU and HPU eager run SDPA with dropout through the math path.
+        return _sdpa_math_dropout(query, key, value, dropout_p)
+    return F.scaled_dot_product_attention(
+        query, key, value, attn_mask=attn_mask, dropout_p=dropout_p
+    )
+
+
 class MultiHeadAttention(nn.Module):
     """Multi-head self-attention block.
 

braindecode-1.8.1.dev168570878/braindecode/version.py

@@ -0,0 +1 @@
+__version__ = "1.8.1.dev168570878"

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878/braindecode.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.8.1.dev168512945
+Version: 1.8.1.dev168570878
 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.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/braindecode.egg-info/SOURCES.txt

@@ -50,7 +50,6 @@ braindecode/datautil/serialization.py
 braindecode/datautil/util.py
 braindecode/functional/__init__.py
 braindecode/functional/_real_dft.py
-braindecode/functional/attention.py
 braindecode/functional/functions.py
 braindecode/functional/initialization.py
 braindecode/models/__init__.py

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/docs/api.rst

@@ -485,7 +485,6 @@ The functional module contains various functions that can be used like functiona
      rotate_half
      rotate_pairs
      safe_log
-     scaled_dot_product_attention
      sinusoidal_positional_encoding
      spectral_input
      square

{braindecode-1.8.1.dev168512945 → braindecode-1.8.1.dev168570878}/docs/whats_new.rst

@@ -28,9 +28,6 @@ Current 1.8.1 (2026-08-31)
 Enhancements
 ============
 
-- CI: ``test_core_dependencies_are_imported`` parsed the whole package once per installed distribution (215 s, the last test to finish); it parses once. The test job stops after 60 minutes instead of 6 hours when a worker hangs (:gh:`1471` by `Bruno Aristimunha`_).
-- :func:`braindecode.functional.sinusoidal_positional_encoding` takes a ``base`` (longest wavelength); PopulationTransformer and SignalJEPA build their tables with it, unchanged (:gh:`1458` by `Bruno Aristimunha`_).
-- Add :func:`braindecode.functional.scaled_dot_product_attention` (the shared SDPA helper, now public, with ``scale`` and an ``exact`` op-order mode); ATCNet, AttnSleep, EEGDINO, TCFormer and TMSA-Net compute their attention through it with outputs, losses and gradients unchanged bit for bit (:gh:`1457` by `Bruno Aristimunha`_).
 - Add :func:`braindecode.functional.rotary_frequencies`, :func:`~braindecode.functional.rotary_angles`, :func:`~braindecode.functional.apply_rotary` and :func:`~braindecode.functional.rotate_half`; BaRISTA, BrainOmni, DIVER-1, EpiNT, LUNA and TCFormer build and apply their rotary embeddings with them (parameters and outputs unchanged), and EEGPT loses a rotary code path it never used (:gh:`1456` by `Bruno Aristimunha`_).
 - Add a test that lists every private copy of a standard block (attention, transformer block, MLP, positional encoding, norm, patch embedding, residual) in :mod:`braindecode.models`, so a new copy fails CI unless it is listed with a reason; the new-model checklist asks for the shared blocks (:gh:`1455` by `Bruno Aristimunha`_).
 - Add a test that lists every private copy of a standard block (attention, transformer block, MLP, positional encoding, norm, patch embedding, residual) in :mod:`braindecode.models`, so a new copy fails CI unless it is listed with a reason; the new-model checklist asks for the shared blocks (:gh:`1455` by `Bruno Aristimunha`_).

braindecode-1.8.1.dev168512945/braindecode/functional/attention.py

@@ -1,152 +0,0 @@
-"""Scaled dot-product attention with the device-specific paths in one place."""
-
-# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
-#
-# License: BSD (3-clause)
-
-import math
-from typing import Optional
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-
-
-@torch.jit.unused
-def _hpu_eager_fused_sdpa(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    scale: Optional[float] = None,
-    fused_in_lazy: bool = False,
-) -> Optional[Tensor]:
-    # Lazy mode: None, the graph compiler already fuses torch's SDPA there
-    # (unless the caller measured the fused kernel faster for its shapes).
-    from habana_frameworks.torch.utils.internal import is_lazy
-
-    if is_lazy() and not fused_in_lazy:
-        return None
-    from habana_frameworks.torch.hpex.kernels import FusedSDPA
-
-    # Strided q/k/v views (heads split off a projection) give wrong outputs and
-    # gradients at some shapes; contiguous copies give the right result.
-    query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
-    return FusedSDPA.apply(query, key, value, None, 0.0, False, scale, "None", False)
-
-
-@torch.jit.unused
-def _hpu_lazy() -> bool:
-    from habana_frameworks.torch.utils.internal import is_lazy
-
-    return is_lazy()
-
-
-def _sdpa_math_dropout(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    dropout_p: float,
-    scale: Optional[float] = None,
-) -> Tensor:
-    # torch's math path for SDPA with dropout, op for op, without the pass of
-    # ``_safe_softmax`` that zeroes fully masked rows (none without a mask):
-    # same values, same dropout draws.
-    if scale is None:
-        scale = 1.0 / math.sqrt(query.size(-1))
-    half = math.sqrt(scale)
-    attn = torch.matmul(query * half, key.transpose(-2, -1) * half)
-    attn = F.dropout(attn.softmax(-1), dropout_p, True)
-    return torch.matmul(attn, value)
-
-
-def scaled_dot_product_attention(
-    query: Tensor,
-    key: Tensor,
-    value: Tensor,
-    attn_mask: Optional[Tensor] = None,
-    dropout_p: float = 0.0,
-    scale: Optional[float] = None,
-    fused_in_lazy: bool = False,
-    exact: bool = False,
-) -> Tensor:
-    """``F.scaled_dot_product_attention``; Gaudi's fused kernel on HPU.
-
-    Parameters
-    ----------
-    query, key, value : torch.Tensor
-        Shape ``(..., seq, head_dim)``; ``key``/``value`` may have another
-        sequence length.
-    attn_mask : torch.Tensor, optional
-        Boolean (``True`` takes part) or additive mask, as in torch.
-    dropout_p : float, default=0.0
-        Dropout on the attention weights; pass ``0.0`` in eval mode.
-    scale : float, optional
-        Score scale; ``None`` is ``1 / sqrt(head_dim)``.
-    fused_in_lazy : bool, default=False
-        Use Gaudi's fused kernel in HPU lazy mode too.
-    exact : bool, default=False
-        Run the textbook op order instead of torch's kernels:
-        ``softmax((q @ k^T) / sqrt(d)) @ v``, or ``softmax((q @ k^T) * scale) @ v``
-        when ``scale`` is given, with dropout on the weights. Bit-identical to
-        a hand-written attention; torch's kernels round differently (its math
-        path scales ``q`` and ``k`` by ``sqrt(scale)`` each before the matmul,
-        its fused CPU kernel runs a blocked softmax), about 1e-6 apart.
-
-    Returns
-    -------
-    torch.Tensor
-        Shape ``(..., seq, head_dim)``.
-
-    In HPU eager mode, torch's SDPA runs op by op (the full
-    ``(batch, heads, seq, seq)`` scores in device memory), 2-3.5x slower than
-    the bridge's fused kernel on long sequences. The fused kernel is used
-    without a mask or dropout (its dropout would draw other masks). In lazy
-    mode only with ``fused_in_lazy=True``: the graph compiler fuses torch's
-    SDPA too, faster on long sequences, slower on short ones.
-    """
-    if exact:
-        scores = torch.matmul(query, key.transpose(-2, -1))
-        if scale is None:
-            scores = scores / math.sqrt(query.size(-1))
-        else:
-            scores = scores * scale
-        if attn_mask is not None:
-            if attn_mask.dtype == torch.bool:
-                scores = scores.masked_fill(~attn_mask, float("-inf"))
-            else:
-                scores = scores + attn_mask
-        weights = scores.softmax(dim=-1)
-        if dropout_p > 0.0:
-            weights = F.dropout(weights, dropout_p, True)
-        return torch.matmul(weights, value)
-    if (
-        not torch.jit.is_scripting()
-        and query.device.type == "hpu"
-        and attn_mask is None
-        and dropout_p == 0.0
-        and (query.dtype == torch.float32 or query.dtype == torch.bfloat16)
-    ):
-        out = _hpu_eager_fused_sdpa(
-            query, key, value, scale=scale, fused_in_lazy=fused_in_lazy
-        )
-        if out is not None:
-            return out
-    if (
-        not torch.jit.is_scripting()
-        and dropout_p > 0.0
-        and attn_mask is None
-        and (query.dtype == torch.float32 or query.dtype == torch.float64)
-        and (
-            query.device.type == "cpu"
-            or (query.device.type == "hpu" and not _hpu_lazy())
-        )
-    ):
-        # CPU and HPU eager run SDPA with dropout through the math path.
-        return _sdpa_math_dropout(query, key, value, dropout_p, scale)
-    if scale is None:  # torch 2.0's signature has no ``scale``
-        return F.scaled_dot_product_attention(
-            query, key, value, attn_mask=attn_mask, dropout_p=dropout_p
-        )
-    return F.scaled_dot_product_attention(
-        query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, scale=scale
-    )

braindecode-1.8.1.dev168512945/braindecode/version.py

@@ -1 +0,0 @@
-__version__ = "1.8.1.dev168512945"