From ac0a78fe22b12c977216e86a7c5978d94c1cbf38 Mon Sep 17 00:00:00 2001 From: Vaibhav Balani Date: Tue, 14 Jul 2026 04:18:05 +0530 Subject: [PATCH 1/2] ENH: Add array_api support for generalized eigenvalue decomposition (TPU/GPU) --- mne/decoding/_ged.py | 30 +++++++++++++++++++++++++++--- mne/decoding/base.py | 9 ++++++++- 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/mne/decoding/_ged.py b/mne/decoding/_ged.py index ed6915b6674..d74a19f7ad0 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..fc0504fae3d 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]) From 029f46cc0cb5af8e6cc7daaacbbf94f0122b6e22 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 Jul 2026 23:15:44 +0000 Subject: [PATCH 2/2] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- mne/decoding/_ged.py | 8 ++++---- mne/decoding/base.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/mne/decoding/_ged.py b/mne/decoding/_ged.py index d74a19f7ad0..4df5094f8c4 100644 --- a/mne/decoding/_ged.py +++ b/mne/decoding/_ged.py @@ -47,7 +47,7 @@ def _smart_ged(S, R, restr_mat=None, R_func=None): import numpy as xp if restr_mat is None: - if xp.__name__ == 'numpy': + if xp.__name__ == "numpy": evals, evecs = scipy.linalg.eigh(S, R) return evals, evecs else: @@ -63,8 +63,8 @@ def _smart_ged(S, R, restr_mat=None, R_func=None): R_restr = restr_mat @ R @ restr_mat.T if R_func is not None: R_restr = R_func([S_restr, R_restr]) - - if xp.__name__ == 'numpy': + + if xp.__name__ == "numpy": evals, evecs_restr = scipy.linalg.eigh(S_restr, R_restr) else: L = xp.linalg.cholesky(R_restr) @@ -72,7 +72,7 @@ def _smart_ged(S, R, restr_mat=None, R_func=None): 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 fc0504fae3d..5bcebf4abdc 100644 --- a/mne/decoding/base.py +++ b/mne/decoding/base.py @@ -210,14 +210,14 @@ def fit(self, X, y=None): ) self._validate_ged_params() covs, C_ref, info, rank, kwargs = self.cov_callable(X, y) - + # 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])