diff --git a/mne/decoding/_ged.py b/mne/decoding/_ged.py index ed6915b6674..4df5094f8c4 100644 --- a/mne/decoding/_ged.py +++ b/mne/decoding/_ged.py @@ -40,15 +40,39 @@ def _smart_ged(S, R, restr_mat=None, R_func=None): to the original space. The g-eigenvectors matrix is of shape (n_chs, r). If callable R_func is provided the GED will be performed on (S, R_func(S,R)) """ + # Array API Standard support for hardware acceleration (TPU/GPU) + if hasattr(S, "__array_namespace__"): + xp = S.__array_namespace__() + else: + import numpy as xp + if restr_mat is None: - evals, evecs = scipy.linalg.eigh(S, R) - return evals, evecs + if xp.__name__ == "numpy": + evals, evecs = scipy.linalg.eigh(S, R) + return evals, evecs + else: + # Hardware-accelerated Cholesky path since eigh(A,B) is not in Array API + L = xp.linalg.cholesky(R) + L_inv = xp.linalg.inv(L) + C = L_inv @ S @ L_inv.T + evals, evecs_C = xp.linalg.eigh(C) + evecs = L_inv.T @ evecs_C + return evals, evecs S_restr = restr_mat @ S @ restr_mat.T R_restr = restr_mat @ R @ restr_mat.T if R_func is not None: R_restr = R_func([S_restr, R_restr]) - evals, evecs_restr = scipy.linalg.eigh(S_restr, R_restr) + + if xp.__name__ == "numpy": + evals, evecs_restr = scipy.linalg.eigh(S_restr, R_restr) + else: + L = xp.linalg.cholesky(R_restr) + L_inv = xp.linalg.inv(L) + C = L_inv @ S_restr @ L_inv.T + evals, evecs_C = xp.linalg.eigh(C) + evecs_restr = L_inv.T @ evecs_C + evecs = restr_mat.T @ evecs_restr return evals, evecs diff --git a/mne/decoding/base.py b/mne/decoding/base.py index 75ec826b711..5bcebf4abdc 100644 --- a/mne/decoding/base.py +++ b/mne/decoding/base.py @@ -210,7 +210,14 @@ def fit(self, X, y=None): ) self._validate_ged_params() covs, C_ref, info, rank, kwargs = self.cov_callable(X, y) - covs = np.stack(covs) + + # Array API Standard support to preserve hardware acceleration tensors + if len(covs) > 0 and hasattr(covs[0], "__array_namespace__"): + xp = covs[0].__array_namespace__() + covs = xp.stack(covs) + else: + covs = np.stack(covs) + self._validate_covariances(covs) if C_ref is not None: self._validate_covariances([C_ref])