ENH: Add array_api support for generalized eigenvalue decomposition (…#14053
Open
V-bot5123 wants to merge 3 commits into
Open
ENH: Add array_api support for generalized eigenvalue decomposition (…#14053V-bot5123 wants to merge 3 commits into
V-bot5123 wants to merge 3 commits into
Conversation
Member
|
Any AI tools used here (see disclosure policy here)? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
As discussed in the issue thread, this PR introduces hardware acceleration (TPU/GPU) for the _smart_ged solver used by CSP, SPoC, and SSD.
Per @larsoner's suggestion, rather than hardcoding PyTorch/XLA or CuPy, this implementation leverages the Array API Standard (array_namespace). This dynamically routes the computation to whatever hardware the user provides, while maintaining a seamless zero-dependency fallback to scipy for standard NumPy users.
🧠 Logical & Architectural Changes
Array API Detection (_ged.py): The _smart_ged function now detects if the incoming matrices S and R support the Array API standard. If they do (e.g. PyTorch Tensors, JAX arrays), we use their native hardware-accelerated linear algebra operations.
Tensor Preservation (base.py): Modified _GEDTransformer.fit to use xp.stack(covs) instead of np.stack(covs). Previously, if a user passed hardware-accelerated tensors into the cov_callable, np.stack would force them back onto the CPU as NumPy arrays before they ever reached the eigensolver. They are now preserved on the device.
📐 Mathematical Changes
Because scipy.linalg.eigh(S, R) natively supports the Generalized Eigenvalue Problem (S * v = lambda * R * v), but most hardware accelerators (like torch.linalg.eigh) only support the Standard Eigenvalue Problem (C * y = lambda * y), we mathematically rewrote the solver to run on hardware accelerators via a Cholesky decomposition.
When a hardware tensor is detected, the code performs the following sequence to solve it natively on the accelerator:
Cholesky Decomposition: We decompose the positive-definite matrix R into a lower triangular matrix L such that R = L * L.T
Standardization: We construct a standard symmetric matrix C = L_inv * S * L_inv.T
Hardware Eigensolve: We solve C * y = lambda * y natively on the TPU/GPU.
Eigenvector Recovery: We transform back to original space using v = L_inv.T * y (prompted Gemini till want satisfied with the formalized way of presenting the text)