braindecode 1.3.2.dev179983996__tar.gz → 1.3.2.dev182544824__tar.gz

{braindecode-1.3.2.dev179983996/braindecode.egg-info → braindecode-1.3.2.dev182544824}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.2.dev179983996
+Version: 1.3.2.dev182544824
 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>
@@ -22,7 +22,7 @@ Requires-Python: >=3.11
 Description-Content-Type: text/x-rst
 License-File: LICENSE.txt
 License-File: NOTICE.txt
-Requires-Dist: torch>=2.2
+Requires-Dist: torch>=2.0
 Requires-Dist: torchaudio>=2.0
 Requires-Dist: mne>=1.11.0
 Requires-Dist: pandas<3.0.0
@@ -35,7 +35,7 @@ Requires-Dist: wfdb
 Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
 Requires-Dist: rotary_embedding_torch
-Requires-Dist: pandas<3.0.0
+Requires-Dist: pydantic>=2.0
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.4.3; extra == "moabb"
 Provides-Extra: eegprep

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/datasets/bids/hub.py

@@ -845,11 +845,16 @@ class HubDatasetMixin:
 
                 # Convert to MNE objects and create dataset
                 info = Info.from_json_dict(info_dict)
+                targets = metadata["target"].values
+                if np.issubdtype(targets.dtype, np.integer):
+                    event_ids = targets
+                else:
+                    event_ids = np.ones(len(metadata), dtype=int)
                 events = np.column_stack(
                     [
-                        metadata["i_start_in_trial"].values,
+                        metadata["i_start_in_trial"].values.astype(int),
                         np.zeros(len(metadata), dtype=int),
-                        metadata["target"].values,
+                        event_ids,
                     ]
                 )
                 epochs = mne.EpochsArray(data, info, events=events, metadata=metadata)

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/__init__.py

@@ -55,6 +55,7 @@ from .tsinception import TSception
 from .usleep import USleep
 from .util import (
     _init_models_dict,
+    build_model_config,
     extract_channel_locations_from_chs_info,
     models_mandatory_parameters,
 )
@@ -118,6 +119,7 @@ __all__ = [
     "TIDNet",
     "TSception",
     "USleep",
+    "build_model_config",
     "_init_models_dict",
     "models_mandatory_parameters",
 ]

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/base.py

@@ -17,6 +17,13 @@ from docstring_inheritance import NumpyDocstringInheritanceInitMeta
 from mne.utils import _soft_import
 from torchinfo import ModelStatistics, summary
 
+from braindecode.models.util import (
+    _EEG_PARAMS,
+    _IMPORT_ADAPTER,
+    build_model_config,
+    resolve_type_kwargs,
+    track_model_init_kwargs,
+)
 from braindecode.version import __version__
 
 huggingface_hub = _soft_import(
@@ -86,61 +93,69 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
     -----
     If some input signal-related parameters are not specified,
     there will be an attempt to infer them from the other parameters.
+    """
 
+    #: Template for model-specific Hub integration notes appended by
+    #: ``__init_subclass__``.  ``{name}`` is replaced with the concrete
+    #: model class name.
+    _HUB_NOTES_TEMPLATE = """
     .. rubric:: Hugging Face Hub integration
 
     When the optional ``huggingface_hub`` package is installed, all models
     automatically gain the ability to be pushed to and loaded from the
     Hugging Face Hub. Install with::
 
-        pip install braindecode[hug]
+        pip install braindecode[hub]
 
     **Pushing a model to the Hub:**
 
-    .. code-block:: python
+    .. code-block::
 
-        from braindecode.models import EEGNetv4
+        from braindecode.models import {name}
 
         # Train your model
-        model = EEGNetv4(n_chans=22, n_outputs=4, n_times=1000)
+        model = {name}(n_chans=22, n_outputs=4, n_times=1000)
         # ... training code ...
 
         # Push to the Hub
         model.push_to_hub(
-            repo_id="username/my-eegnet-model", commit_message="Initial model upload"
+            repo_id="username/my-{name_lower}-model",
+            commit_message="Initial model upload",
         )
 
     **Loading a model from the Hub:**
 
-    .. code-block:: python
+    .. code-block::
 
-        from braindecode.models import EEGNetv4
+        from braindecode.models import {name}
 
         # Load pretrained model
-        model = EEGNetv4.from_pretrained("username/my-eegnet-model")
+        model = {name}.from_pretrained("username/my-{name_lower}-model")
 
     The integration automatically handles EEG-specific parameters (n_chans,
     n_times, sfreq, chs_info, etc.) by saving them in a config file alongside
     the model weights. This ensures that loaded models are correctly configured
     for their original data specifications.
 
-    .. important::
-        Currently, only EEG-specific parameters (n_outputs, n_chans, n_times,
-        input_window_seconds, sfreq, chs_info) are saved to the Hub. Model-specific
-        parameters (e.g., dropout rates, activation functions, number of filters)
-        are not preserved and will use their default values when loading from the Hub.
-
-        To use non-default model parameters, specify them explicitly when calling
-        :func:`from_pretrained()`::
-
-            model = EEGNet.from_pretrained("user/model", dropout=0.3, activation='relu')
-
-        Full parameter serialization will be addressed in a future update.
+    All model parameters (both EEG-specific and model-specific such as
+    dropout rates, activation functions, number of filters) are automatically
+    saved to the Hub and restored when loading.
     """
 
     def __init_subclass__(cls, **kwargs):
+        # Append model-specific Hub integration notes to the docstring.
+        # This runs after the metaclass, so we concatenate rather than
+        # override any existing Notes section in the subclass.
+        if cls.__doc__ is not None:
+            hub_notes = cls._HUB_NOTES_TEMPLATE.format(
+                name=cls.__name__,
+                name_lower=cls.__name__.lower(),
+            )
+            cls.__doc__ = cls.__doc__.rstrip() + "\n" + hub_notes
+
         if not HAS_HF_HUB:
             super().__init_subclass__(**kwargs)
+            track_model_init_kwargs(cls)
             return
 
         base_tags = ["braindecode", cls.__name__]
@@ -157,6 +172,19 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
         repo_url = kwargs.pop("repo_url", "https://braindecode.org")
         library_name = kwargs.pop("library_name", "braindecode")
         license = kwargs.pop("license", "bsd-3-clause")
+
+        # Register a coder so that type[nn.Module] parameters
+        # (e.g. activation=nn.ELU) are serialized as importable
+        # strings in config.json and decoded back on load.
+        coders = kwargs.pop("coders", None) or {}
+        coders.setdefault(
+            type,
+            (
+                lambda t: f"{t.__module__}.{t.__qualname__}",
+                lambda data: _IMPORT_ADAPTER.validate_python(data),
+            ),
+        )
+
         # TODO: model_card_template can be added in the future for custom model cards
         super().__init_subclass__(
             tags=tags,
@@ -164,8 +192,10 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
             repo_url=repo_url,
             library_name=library_name,
             license=license,
+            coders=coders,
             **kwargs,
         )
+        track_model_init_kwargs(cls)
 
     def __init__(
         self,
@@ -205,6 +235,17 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
         self._n_times = n_times  # type: ignore[assignment]
         self._sfreq = sfreq  # type: ignore[assignment]
 
+        # Back-fill instance attributes from _hub_mixin_config for any
+        # params the subclass didn't store on self.  Skip EEG params
+        # (stored as self._*) and descriptors (properties).
+        for key, val in getattr(self, "_hub_mixin_config", {}).items():
+            if key in _EEG_PARAMS:
+                continue
+            if hasattr(getattr(type(self), key, None), "__get__"):
+                continue
+            if not hasattr(self, key):
+                setattr(self, key, val)
+
         super().__init__()
 
     @property
@@ -315,6 +356,76 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
                     raise ValueError(msg) from exc
                 raise exc
 
+    def get_config(self) -> dict:
+        """Return a JSON-serializable dict of all ``__init__`` parameters.
+
+        The returned dictionary can be saved to a JSON file and later
+        used with :meth:`from_config` to reconstruct the model (without
+        weights).  It is also used internally by :meth:`push_to_hub` to
+        persist the full model configuration.
+
+        Returns
+        -------
+        dict
+            All ``__init__`` parameters, JSON-serializable.
+            ``type[nn.Module]`` parameters (e.g. ``activation``) are
+            encoded as importable dotted-path strings.
+
+        Examples
+        --------
+        >>> import json
+        >>> from braindecode.models import EEGNet
+        >>> model = EEGNet(n_chans=22, n_times=1000, n_outputs=4, F1=16)
+        >>> config = model.get_config()
+        >>> config["F1"]
+        16
+        >>> # Save to disk
+        >>> with open("config.json", "w") as f:
+        ...     json.dump(config, f)
+
+        .. versionadded:: 1.4
+        """
+        return build_model_config(self)
+
+    @classmethod
+    def from_config(cls, config: dict) -> "EEGModuleMixin":
+        """Create a model instance from a configuration dict.
+
+        This is the inverse of :meth:`get_config`.  Weights are **not**
+        loaded -- use :meth:`from_pretrained` for that.
+
+        Parameters
+        ----------
+        config : dict
+            Configuration dict as returned by :meth:`get_config`.
+
+        Returns
+        -------
+        EEGModuleMixin
+            A new model instance.
+
+        Examples
+        --------
+        >>> import json
+        >>> from braindecode.models import EEGNet
+        >>> model = EEGNet(n_chans=22, n_times=1000, n_outputs=4, F1=16)
+        >>> config = model.get_config()
+        >>> # Reconstruct (without weights)
+        >>> model2 = EEGNet.from_config(config)
+        >>> model2.F1
+        16
+        >>> # Or from a JSON file
+        >>> with open("config.json") as f:
+        ...     config = json.load(f)
+        >>> model3 = EEGNet.from_config(config)
+
+        .. versionadded:: 1.4
+        """
+        config = dict(config)  # shallow copy
+        config.pop("braindecode_version", None)
+        resolve_type_kwargs(cls, config)
+        return cls(**config)
+
     mapping: Optional[Dict[str, str]] = None
 
     def load_state_dict(self, state_dict, *args, **kwargs):
@@ -419,54 +530,22 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
 
     @staticmethod
     def _serialize_chs_info(chs_info):
-        """
-        Serialize MNE channel info to JSON-compatible format.
-
-        Parameters
-        ----------
-        chs_info : list of dict or None
-            Channel information from MNE Info object.
-
-        Returns
-        -------
-        list of dict or None
-            Serialized channel information that can be saved to JSON.
-        """
+        """Serialize MNE channel info (``info["chs"]``) to JSON-compatible dicts."""
         if chs_info is None:
             return None
-
+        _INT_FIELDS = ("kind", "coil_type", "unit")
+        _FLOAT_FIELDS = ("cal", "range")
         serialized = []
         for ch in chs_info:
-            # Extract serializable fields from MNE channel info
-            ch_dict = {
-                "ch_name": ch.get("ch_name", ""),
-            }
-
-            # Handle kind field - can be either string or integer
-            kind_val = ch.get("kind")
-            if kind_val is not None:
-                ch_dict["kind"] = (
-                    kind_val if isinstance(kind_val, str) else int(kind_val)
-                )
-
-            # Add numeric fields with safe conversion
-            coil_type = ch.get("coil_type")
-            if coil_type is not None:
-                ch_dict["coil_type"] = int(coil_type)
-
-            unit = ch.get("unit")
-            if unit is not None:
-                ch_dict["unit"] = int(unit)
-
-            cal = ch.get("cal")
-            if cal is not None:
-                ch_dict["cal"] = float(cal)
-
-            range_val = ch.get("range")
-            if range_val is not None:
-                ch_dict["range"] = float(range_val)
-
-            # Serialize location array if present
+            ch_dict = {"ch_name": ch.get("ch_name", "")}
+            for key in _INT_FIELDS:
+                val = ch.get(key)
+                if val is not None:
+                    ch_dict[key] = val if isinstance(val, str) else int(val)
+            for key in _FLOAT_FIELDS:
+                val = ch.get(key)
+                if val is not None:
+                    ch_dict[key] = float(val)
             if "loc" in ch and ch["loc"] is not None:
                 ch_dict["loc"] = (
                     ch["loc"].tolist()
@@ -474,35 +553,19 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
                     else list(ch["loc"])
                 )
             serialized.append(ch_dict)
-
         return serialized
 
     @staticmethod
     def _deserialize_chs_info(chs_info_dict):
-        """
-        Deserialize channel info from JSON-compatible format to MNE-like structure.
-
-        Parameters
-        ----------
-        chs_info_dict : list of dict or None
-            Serialized channel information.
-
-        Returns
-        -------
-        list of dict or None
-            Deserialized channel information compatible with MNE.
-        """
+        """Deserialize JSON channel dicts back to MNE-compatible format."""
         if chs_info_dict is None:
             return None
-
         deserialized = []
         for ch_dict in chs_info_dict:
             ch = ch_dict.copy()
-            # Convert location back to numpy array if present
             if "loc" in ch and ch["loc"] is not None:
                 ch["loc"] = np.array(ch["loc"])
             deserialized.append(ch)
-
         return deserialized
 
     def _save_pretrained(self, save_directory):
@@ -522,16 +585,8 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
 
         save_directory = Path(save_directory)
 
-        # Collect EEG-specific configuration
-        config = {
-            "n_outputs": self._n_outputs,
-            "n_chans": self._n_chans,
-            "n_times": self._n_times,
-            "input_window_seconds": self._input_window_seconds,
-            "sfreq": self._sfreq,
-            "chs_info": self._serialize_chs_info(self._chs_info),
-            "braindecode_version": __version__,
-        }
+        config = build_model_config(self)
+        config["braindecode_version"] = __version__
 
         # Save to config.json
         config_path = save_directory / "config.json"
@@ -570,6 +625,7 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
             **model_kwargs,
         ):
             model_kwargs.pop("braindecode_version", None)
+            resolve_type_kwargs(cls, model_kwargs)
             return super()._from_pretrained(  # type: ignore
                 model_id=model_id,
                 revision=revision,

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/cbramod.py

@@ -207,11 +207,25 @@ class CBraMod(EEGModuleMixin, nn.Module):
 
         self._weights_init()
 
-        self.final_layer = (
-            nn.Identity()
-            if return_encoder_output
-            else nn.Sequential(nn.Flatten(), nn.LazyLinear(self.n_outputs))
-        )
+        try:
+            n_times = self.n_times
+            n_chans = self.n_chans
+        except ValueError:
+            n_times = None
+            n_chans = None
+
+        if return_encoder_output:
+            self.final_layer = nn.Identity()
+        elif n_times is not None and n_chans is not None:
+            n_patch = n_times // patch_size
+            flat_dim = n_chans * n_patch * emb_dim
+            self.final_layer = nn.Sequential(
+                nn.Flatten(), nn.Linear(flat_dim, self.n_outputs)
+            )
+        else:
+            self.final_layer = nn.Sequential(
+                nn.Flatten(), nn.LazyLinear(self.n_outputs)
+            )
 
     def _weights_init(self):
         for m in self.modules():

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/config.py

@@ -4,13 +4,16 @@ from types import UnionType
 from typing import Annotated, Any, Literal, Union, get_args, get_origin
 
 import numpy as np
-from mne.utils import _soft_import
+import pydantic
 from typing_extensions import TypedDict
 
 from braindecode.models.base import EEGModuleMixin
-from braindecode.models.util import SigArgName, models_dict, models_mandatory_parameters
-
-pydantic = _soft_import(name="pydantic", purpose="model configuration", strict=False)
+from braindecode.models.util import (
+    SigArgName,
+    _init_models_dict,
+    models_dict,
+    models_mandatory_parameters,
+)
 
 try:
     from numpydantic import NDArray, Shape
@@ -57,8 +60,12 @@ SIGNAL_ARGS_TYPES = {
 }
 
 
-class BaseBraindecodeModelConfig(pydantic.BaseModel):  # type: ignore
+class BaseBraindecodeModelConfig(pydantic.BaseModel):
+    """Base class for braindecode model pydantic configs."""
+
     def create_instance(self) -> EEGModuleMixin:
+        if self.model_name_ not in models_dict:
+            _init_models_dict()
         model_cls = models_dict[self.model_name_]
         kwargs = self.model_dump(mode="python", exclude={"model_name_"})
         if kwargs.get("n_chans") is not None and kwargs.get("chs_info") is not None:
@@ -90,13 +97,6 @@ def make_model_config(
     type
         A pydantic BaseModel subclass representing the model config.
     """
-    if not pydantic:
-        raise ImportError(
-            "pydantic is required to use make_model_config. "
-            "Please install braindecode[typing]."
-        )
-
-    # ironically, we need to ignore the type here to have the soft dependency.
 
     @pydantic.model_validator(mode="before")
     def validate_signal_params(cls, data: Any):
@@ -207,23 +207,23 @@ def make_model_config(
 # and define __all__ based on generated classes
 __all__ = ["make_model_config"]
 
-if not pydantic:
-    pass
-else:
-    models_configs: list[type[BaseBraindecodeModelConfig]] = []
-    for model_name, req, _ in models_mandatory_parameters:
-        model_cls = models_dict[model_name]
-        model_cfg = make_model_config(model_cls, req)
-        globals()[model_cfg.__name__] = model_cfg
-        __all__.append(model_cfg.__name__)
-        models_configs.append(model_cfg)
-
-    BraindecodeModelConfig = Annotated[  # type: ignore
-        Union[tuple(models_configs)],
-        pydantic.Field(
-            discriminator="model_name_", description="Braindecode model configuration"
-        ),
-    ]
+if not models_dict:
+    _init_models_dict()
+
+models_configs: list[type[BaseBraindecodeModelConfig]] = []
+for model_name, req, _ in models_mandatory_parameters:
+    model_cls = models_dict[model_name]
+    model_cfg = make_model_config(model_cls, req)
+    globals()[model_cfg.__name__] = model_cfg
+    __all__.append(model_cfg.__name__)
+    models_configs.append(model_cfg)
+
+BraindecodeModelConfig = Annotated[  # type: ignore[valid-type]
+    Union[tuple(models_configs)],
+    pydantic.Field(
+        discriminator="model_name_", description="Braindecode model configuration"
+    ),
+]
 
 # # Example usage:
 #

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/eegsym.py

@@ -112,7 +112,7 @@ class EEGSym(EEGModuleMixin, nn.Module):
         corresponding to different frequency characteristics.
 
     Notes
-    ----------
+    -----
     * EEGSym achieved competitive accuracies across five large MI datasets [eegsym2022]_.
     * The model maintained high accuracy using a reduced set of electrodes (8 or 16 channels)
       [eegsym2022]_.

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/luna.py

@@ -7,6 +7,7 @@ Retrieved from https://openreview.net/forum?id=uazfjnFL0G
 
 Original Authors: Berkay Döner, Thorir Mar Ingolfsson
 Braindecode Adaptation: Bruno Aristimunha
+Contributions: Sarthak Tayal <sarthaktayal2@gmail.com>
 
 the LICENSE Of this file is APACHE-2.0.
 """
@@ -82,11 +83,8 @@ class LUNA(EEGModuleMixin, nn.Module):
 
            # Load pre-trained encoder weights
            state_dict = load_file(model_path)
-           # Apply key mapping for pretrained weights
-           mapping = model.mapping.copy()
-           mapping["cross_attn.temparature"] = "cross_attn.temperature"
-           mapped_state_dict = {mapping.get(k, k): v for k, v in state_dict.items()}
-           model.load_state_dict(mapped_state_dict, strict=False)
+           # load_state_dict applies model.mapping automatically
+           model.load_state_dict(state_dict, strict=False)
 
        To push your own trained model to the Hub:
 
@@ -170,6 +168,7 @@ class LUNA(EEGModuleMixin, nn.Module):
             "channel_location_embedder.0.fc2.bias": "channel_location_embedder.fc2.bias",
             "channel_location_embedder.0.norm.weight": "channel_location_embedder.norm.weight",
             "channel_location_embedder.0.norm.bias": "channel_location_embedder.norm.bias",
+            "cross_attn.temparature": "cross_attn.temperature",  # typo in pretrained weights
         }
 
         # Model parameters
@@ -315,8 +314,10 @@ class LUNA(EEGModuleMixin, nn.Module):
         channel_locations_emb = self.channel_location_embedder(channel_locations)
 
         x_tokenized = rearrange(x_masked, "B (C t) D -> (B t) C D", C=num_channels)
-        channel_locations_emb = channel_locations_emb.repeat(
-            num_patches_per_channel, 1, 1
+        channel_locations_emb = (
+            channel_locations_emb.repeat_interleave(  # was repeat(), wrong dim ordering
+                num_patches_per_channel, dim=0
+            )
         )
         x_tokenized = x_tokenized + channel_locations_emb
 

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/patchedtransformer.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import math
+import warnings
 from typing import Optional, Tuple
 
 import torch
@@ -496,9 +497,10 @@ class _TransformerEncoderLayer(nn.Module):
 
         if dim_feedforward is None:
             dim_feedforward = 4 * embed_dim
-            # note: preserve the original behaviour (print) from the provided code
-            print(
-                "dim_feedforward is set to 4*embed_dim, the default in Vaswani et al. (Attention is all you need)"
+            warnings.warn(
+                "dim_feedforward is set to 4*embed_dim, the default in "
+                "Vaswani et al. (Attention is all you need)",
+                stacklevel=2,
             )
 
         self.layer_norm_att = _LayerNorm(embed_dim, bias=bias)

{braindecode-1.3.2.dev179983996 → braindecode-1.3.2.dev182544824}/braindecode/models/reve.py

@@ -16,7 +16,16 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 from torch import nn
-from torch.nn.attention import SDPBackend, sdpa_kernel
+
+# Safe import for older PyTorch versions (Support for Intel-based Macs)
+try:
+    from torch.nn.attention import SDPBackend, sdpa_kernel
+
+    HAS_SDPA = True
+except ImportError:
+    HAS_SDPA = False
+    SDPBackend = None
+    sdpa_kernel = None
 
 from braindecode.models.base import EEGModuleMixin
 
@@ -530,11 +539,22 @@ class FeedForward(nn.Module):
 
 
 class ClassicalAttention(nn.Module):
-    def __init__(self, heads: int, use_sdpa: bool = True):
+    def __init__(self, heads: int, use_sdpa: bool | None = None):
         super().__init__()
-        self.use_sdpa = use_sdpa
         self.heads = heads
 
+        if use_sdpa is None:
+            self.use_sdpa = HAS_SDPA
+        elif use_sdpa is True and not HAS_SDPA:
+            logger.warning(
+                "SDPA (Scaled Dot Product Attention) was requested, but it is not "
+                "available in your current PyTorch version. Falling back to naive "
+                "implementation. Please upgrade to PyTorch >= 2.2 for SDPA support."
+            )
+            self.use_sdpa = False
+        else:
+            self.use_sdpa = use_sdpa
+
     def forward(self, qkv: torch.Tensor) -> torch.Tensor:
         # Split concatenated QKV into separate tensors
         # qkv shape: (batch, seq_len, 3 * heads * head_dim)