mne-tools / mne-tools/mne-python
ENH support the array_API for TRF decoder
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 3.5k
- Forks
- 1.6k
- Avg merge
- 1d 6h
- Merged PRs (30d)
- 100
Description
Describe the new feature or enhancement
On a project with a student, we are relying on mne's TRF model, which is super efficient, thanks for that!
One bottleneck is that we have data embeddings on GPUs, and we would like to stay on GPU for the TRF.
With the current improvement of the array_api support on scikit-learn (see here), Ridge can now be called on GPU with the same object/API.
Would it make sense to adapt mne TRF code to allow using the object with the array API when passed torch tensors for instance?
Describe your proposed implementation
Change the way the lag view are constructed to use the array_api compat layer, and allow pass the data in the original format to the scikit-learn estimator when it is supported.
Describe possible alternatives
N/A
Additional context
No response
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by locating the TRF decoder and the code that constructs lag views, then inspect how the array_api compatibility layer and scikit-learn estimator are used. The change is complete when supported array-API inputs such as torch tensors can remain in their original format through the TRF path and be accepted by the estimator.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch, scikit-learn
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100