braindecode 1.8.1.dev169698056__tar.gz → 1.8.1.dev169869285__tar.gz

{braindecode-1.8.1.dev169698056/braindecode.egg-info → braindecode-1.8.1.dev169869285}/PKG-INFO

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

@@ -1,5 +1,7 @@
+from .attention import scaled_dot_product_attention
 from .functions import (
     _get_gaussian_kernel1d,
+    apply_rotary,
     daubechies_filters,
     drop_path,
     dwt_max_level,
@@ -10,6 +12,9 @@ from .functions import (
     pairwise_iou_1d,
     plv_time,
     prefer_fft_conv,
+    rotary_angles,
+    rotary_frequencies,
+    rotate_half,
     rotate_pairs,
     safe_log,
     sinusoidal_positional_encoding,
@@ -17,10 +22,16 @@ from .functions import (
     square,
     wavelet_decomposition,
 )
-from .initialization import glorot_weight_zero_bias, rescale_parameter
+from .initialization import (
+    glorot_weight_zero_bias,
+    init_weights_,
+    rescale_parameter,
+)
+from .patches import pad_or_crop_to_patches, patch_count
 
 __all__ = [
     "_get_gaussian_kernel1d",
+    "apply_rotary",
     "daubechies_filters",
     "drop_path",
     "dwt_max_level",
@@ -28,15 +39,22 @@ __all__ = [
     "hilbert_freq",
     "identity",
     "iou_1d",
+    "pad_or_crop_to_patches",
     "pairwise_iou_1d",
+    "patch_count",
     "plv_time",
     "prefer_fft_conv",
+    "rotary_angles",
+    "rotary_frequencies",
+    "rotate_half",
     "rotate_pairs",
     "safe_log",
+    "scaled_dot_product_attention",
     "sinusoidal_positional_encoding",
     "spectral_input",
     "square",
     "wavelet_decomposition",
     "glorot_weight_zero_bias",
     "rescale_parameter",
+    "init_weights_",
 ]

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

@@ -0,0 +1,152 @@
+"""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.dev169698056 → braindecode-1.8.1.dev169869285}/braindecode/functional/functions.py

@@ -588,6 +588,7 @@ 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)``.
 
@@ -609,6 +610,9 @@ 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
     -------
@@ -620,7 +624,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(10000.0) / dim_even)
+        * (-math.log(base) / dim_even)
     )
     pe = torch.zeros(n_positions, dim_even, device=device, dtype=dtype)
     pe[:, 0::2] = torch.sin(position * div_term)
@@ -646,3 +650,104 @@ def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
     """
     pairs = x.unflatten(-1, (-1, 2))
     return torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).flatten(-2)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate the two feature halves from ``(a, b)`` to ``(-b, a)``.
+
+    The half-split layout of GPT-NeoX / LLaMA rotary embeddings; see
+    :func:`rotate_pairs` for the interleaved layout.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Tensor of shape ``(..., features)`` with an even final dimension.
+
+    Returns
+    -------
+    torch.Tensor
+        Same shape, device and dtype as ``x``.
+    """
+    first, second = x.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def rotary_frequencies(
+    head_dim: int,
+    base: float = 10000.0,
+    device: Optional[torch.device] = None,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """Inverse frequencies of a rotary embedding, ``1 / base ** (2 i / head_dim)``.
+
+    Parameters
+    ----------
+    head_dim : int
+        Even feature dimension that is rotated.
+    base : float, default=10000.0
+        Base period (``theta``).
+    device : torch.device, optional
+    dtype : torch.dtype, default=torch.float32
+        Precision of the table; the angles inherit it.
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(head_dim // 2,)``.
+    """
+    exponent = torch.arange(0, head_dim, 2, device=device, dtype=dtype) / head_dim
+    return 1.0 / (base**exponent)
+
+
+def rotary_angles(
+    positions: torch.Tensor, inv_freq: torch.Tensor, interleaved: bool = True
+) -> torch.Tensor:
+    """Rotation angles of each position, laid out to match the feature layout.
+
+    Parameters
+    ----------
+    positions : torch.Tensor
+        Shape ``(...,)``; same dtype as ``inv_freq`` for exact angles.
+    inv_freq : torch.Tensor
+        Shape ``(head_dim // 2,)``, from :func:`rotary_frequencies`.
+    interleaved : bool, default=True
+        ``True`` repeats each angle for its adjacent pair (:func:`rotate_pairs`);
+        ``False`` concatenates the angles twice (:func:`rotate_half`).
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(..., head_dim)``; pass ``.cos()`` and ``.sin()`` to
+        :func:`apply_rotary`.
+    """
+    angles = positions.unsqueeze(-1) * inv_freq
+    if interleaved:
+        return angles.repeat_interleave(2, dim=-1)
+    return torch.cat((angles, angles), dim=-1)
+
+
+def apply_rotary(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = True
+) -> torch.Tensor:
+    """Rotate ``x`` by the angles whose cosine and sine are given.
+
+    ``x * cos + rotate(x) * sin``, where ``rotate`` is :func:`rotate_pairs`
+    (``interleaved=True``) or :func:`rotate_half`. ``cos`` and ``sin`` must
+    broadcast to ``x``; casting is left to the caller.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Shape ``(..., head_dim)``.
+    cos, sin : torch.Tensor
+        From :func:`rotary_angles`.
+    interleaved : bool, default=True
+        Feature layout, as in :func:`rotary_angles`.
+
+    Returns
+    -------
+    torch.Tensor
+        Same shape as ``x``.
+    """
+    rotated = rotate_pairs(x) if interleaved else rotate_half(x)
+    return x * cos + rotated * sin

braindecode-1.8.1.dev169869285/braindecode/functional/initialization.py

@@ -0,0 +1,139 @@
+import math
+
+import torch
+from torch import nn
+from torch.nn.utils import parametrize
+
+
+def glorot_weight_zero_bias(model):
+    """Set the weights of BatchNorm layers to one and all biases to zero.
+
+    Despite its name, this function does not apply the Glorot (Xavier)
+    initialization: other weights keep the initialization of their layer
+    (PyTorch's default, or Xavier uniform for the max-norm constrained layers
+    such as :class:`braindecode.modules.Conv2dWithConstraint`). The name is
+    kept for backward compatibility.
+
+    Parameters
+    ----------
+    model : Module
+    """
+    for module in model.modules():
+        # ``weight``/``bias`` can be None (e.g. BatchNorm with affine=False)
+        weight = getattr(module, "weight", None)
+        if isinstance(weight, torch.Tensor) and "BatchNorm" in type(module).__name__:
+            nn.init.constant_(weight, 1)
+        bias = getattr(module, "bias", None)
+        if isinstance(bias, torch.Tensor):
+            nn.init.constant_(bias, 0)
+
+
+def rescale_parameter(param, layer_id):
+    r"""Recaling the l-th transformer layer.
+
+    Rescales the parameter tensor by the inverse square root of the layer id.
+    Made inplace. :math:`\frac{1}{\sqrt{2 \cdot \text{layer\_id}}}` [Beit2022]
+
+    In the labram, this is used to rescale the output matrices
+    (i.e., the last linear projection within each sub-layer) of the
+    self-attention module.
+
+    Parameters
+    ----------
+    param: :class:`torch.Tensor`
+        tensor to be rescaled
+    layer_id: int
+        layer id in the neural network
+
+    References
+    ----------
+    [Beit2022] Hangbo Bao, Li Dong, Songhao Piao, Furu We (2022). BEIT: BERT
+    Pre-Training of Image Transformers.
+    """
+    param.div_(math.sqrt(2.0 * layer_id))
+
+
+def init_weights_(
+    module,
+    weights=None,
+    norms=(nn.LayerNorm,),
+    init_parametrized_original=False,
+):
+    """Initialize one module by layer type; apply to a model with ``model.apply``.
+
+    The first entry of ``weights`` whose type matches ``module`` gives its rule,
+    as in an ``if``/``elif`` chain; modules of a type in ``norms`` get a unit
+    weight and a zero bias; nothing else is touched. The default is the
+    timm/ViT scheme (truncated-normal ``nn.Linear`` weights with std 0.02,
+    unit/zero ``nn.LayerNorm``). Models pass their own rules to reproduce their
+    reference initialization exactly.
+
+    Parameters
+    ----------
+    module : torch.nn.Module
+    weights : dict or None
+        ``{layer type or tuple of types: rule}``. A rule is a function applied
+        to the ``weight`` (the ``bias`` is then set to zero), or a pair
+        ``(weight_fn, bias_fn)`` where either may be ``None`` to leave that
+        tensor as it is. For a layer without a ``weight`` attribute (e.g.
+        ``nn.LSTM``) the functions apply to every parameter whose name contains
+        ``"weight"`` / ``"bias"``. ``None`` uses
+        ``{nn.Linear: partial(nn.init.trunc_normal_, std=0.02)}``.
+    norms : tuple of type
+        Normalization layer types set to weight 1 and bias 0 (missing affine
+        parameters are skipped). ``()`` leaves normalization layers alone.
+    init_parametrized_original : bool
+        For a layer whose ``weight`` is parametrized (e.g. max-norm), initialize
+        the unconstrained ``parametrizations.weight.original`` instead of the
+        computed ``weight``.
+
+    Examples
+    --------
+    >>> from functools import partial
+    >>> from torch import nn
+    >>> model = nn.Sequential(nn.Linear(4, 4), nn.LayerNorm(4))
+    >>> _ = model.apply(init_weights_)  # timm/ViT default
+    >>> kaiming = partial(nn.init.kaiming_normal_, mode="fan_out")
+    >>> rule = {nn.Linear: (kaiming, None)}  # keep the default bias
+    >>> _ = model.apply(partial(init_weights_, weights=rule, norms=()))
+    """
+    if weights is None:
+        weights = {nn.Linear: _trunc_normal_002}
+    for types, rule in weights.items():
+        if isinstance(module, types):
+            weight_fn, bias_fn = (
+                rule if isinstance(rule, tuple) else (rule, nn.init.zeros_)
+            )
+            if not hasattr(module, "weight"):  # e.g. nn.LSTM: weight_ih_l0, ...
+                for name, param in module.named_parameters(recurse=False):
+                    fn = (
+                        weight_fn
+                        if "weight" in name
+                        else bias_fn
+                        if "bias" in name
+                        else None
+                    )
+                    if fn is not None:
+                        fn(param)
+                return
+            if isinstance(module.weight, nn.parameter.UninitializedParameter):
+                return  # lazy layer: initialized when it gets its shape
+            if weight_fn is not None:
+                weight = module.weight
+                if init_parametrized_original and parametrize.is_parametrized(
+                    module, "weight"
+                ):
+                    weight = module.parametrizations.weight.original
+                weight_fn(weight)
+            if bias_fn is not None and getattr(module, "bias", None) is not None:
+                bias_fn(module.bias)
+            return
+    if norms and isinstance(module, norms):
+        if getattr(module, "weight", None) is not None:
+            nn.init.ones_(module.weight)
+        if getattr(module, "bias", None) is not None:
+            nn.init.zeros_(module.bias)
+
+
+def _trunc_normal_002(tensor):
+    return nn.init.trunc_normal_(tensor, std=0.02)

braindecode-1.8.1.dev169869285/braindecode/functional/patches.py

@@ -0,0 +1,105 @@
+"""Cutting a time axis into patches: how many, and padding or cropping to fit."""
+
+# Authors: Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+
+import torch
+import torch.nn.functional as F
+
+_ON_NON_DIVISIBLE = ("pad", "crop", "error")
+
+
+def _check_on_non_divisible(on_non_divisible: str) -> None:
+    """Raise a ``ValueError`` for an unknown ``on_non_divisible`` value."""
+    if on_non_divisible not in _ON_NON_DIVISIBLE:
+        raise ValueError(
+            "on_non_divisible must be 'pad', 'crop', or 'error', "
+            f"got {on_non_divisible!r}."
+        )
+
+
+def patch_count(
+    n_times: int, patch_size: int, on_non_divisible: str = "pad", step: int = 0
+) -> int:
+    """Number of patches cut from ``n_times`` samples.
+
+    Parameters
+    ----------
+    n_times : int
+        Length of the axis that is patched.
+    patch_size : int
+        Samples per patch.
+    on_non_divisible : {"pad", "crop", "error"}, default="pad"
+        A last, incomplete patch is zero-padded (``"pad"``), dropped
+        (``"crop"``) or raises a ``ValueError``.
+    step : int, default=0
+        Distance between patch starts; ``0`` means ``patch_size``
+        (non-overlapping).
+
+    Returns
+    -------
+    int
+        Number of patches.
+    """
+    if step <= 0:
+        step = patch_size
+    if n_times < 1:
+        raise ValueError(f"n_times must be positive, got {n_times}.")
+    if n_times < patch_size:
+        if on_non_divisible == "pad":
+            return 1
+        if on_non_divisible == "crop":
+            raise ValueError(
+                f"Need at least one full patch of {patch_size} samples, got {n_times}."
+            )
+    elif (n_times - patch_size) % step == 0:
+        return (n_times - patch_size) // step + 1
+    elif on_non_divisible != "error":
+        return (n_times - patch_size) // step + 1 + (on_non_divisible == "pad")
+    raise ValueError(
+        f"Time dimension ({n_times}) is not divisible by patch_size ({patch_size})"
+        + ("." if step == patch_size else f" with patch step {step}.")
+    )
+
+
+def pad_or_crop_to_patches(
+    x: torch.Tensor,
+    patch_size: int,
+    on_non_divisible: str = "pad",
+    step: int = 0,
+    min_patches: int = 1,
+    center: bool = False,
+) -> torch.Tensor:
+    """Right zero-pad or crop the last axis so the patches tile it.
+
+    A tiled axis is returned as is. An axis short of ``min_patches`` patches is
+    zero-padded up to that many, with half the missing ones (whole steps)
+    before it if ``center``.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Shape ``(..., n_times)``.
+    patch_size, on_non_divisible, step
+        As in :func:`patch_count`.
+    min_patches : int, default=1
+        Least number of patches on the returned axis.
+    center : bool, default=False
+        Pad half the missing patches on the left.
+
+    Returns
+    -------
+    torch.Tensor
+        Shape ``(..., patch_size + (n_patches - 1) * step)``.
+    """
+    n_patches = patch_count(x.shape[-1], patch_size, on_non_divisible, step)
+    step = step if step > 0 else patch_size
+    left = 0
+    if n_patches < min_patches:
+        left = (min_patches - n_patches) // 2 * step if center else 0
+        n_patches = min_patches
+    target = patch_size + (n_patches - 1) * step
+    if target > x.shape[-1]:
+        return F.pad(x, (left, target - x.shape[-1] - left))
+    return x[..., :target]

{braindecode-1.8.1.dev169698056 → braindecode-1.8.1.dev169869285}/braindecode/models/atcnet.py

@@ -1,7 +1,6 @@
 # Authors: Cedric Rommel <cedric.rommel@inria.fr>
 #
 # License: BSD (3-clause)
-import math
 from typing import Optional
 
 import torch
@@ -10,6 +9,7 @@ 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,
@@ -826,6 +826,12 @@ class _TCNResidualBlock(nn.Module):
         # (reshaping_conv may be nn.Identity, which the helper skips).
         _register_max_norm(conv_max_norm, self.conv1, self.conv2, self.reshaping_conv)
 
+    def reset_parameters(self) -> None:
+        """Initialize this block's weights (the constructor initializes each
+        convolution as it builds it)."""
+        nn.init.kaiming_uniform_(self.conv1.weight)
+        nn.init.kaiming_uniform_(self.conv2.weight)
+
     def forward(self, X):
         # Dimension: (batch_size, F2, Tw)
         # ----- Double dilated convolutions -----
@@ -928,20 +934,10 @@ 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')
 
-        # 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.
+        # Multihead output (batch_size, seq_len, dim): softmax(Q K^T / sqrt(d)) V
+        # per head, heads merged back along the feature axis.
         H = torch.cat(
-            (W.bmm(V_)).split(  # (B', S, S)  # (B', S, D')
-                batch_size, 0
-            ),  # [(B, S, D')] * num_heads
+            scaled_dot_product_attention(Q_, K_, V_, exact=True).split(batch_size, 0),
             -1,
         )  # (B, S, D)