mne-tools / mne-tools/mne-python

ENH support the array_API for TRF decoder

Open
#13,876 4 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

ENH
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.