Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Dec 3, 2024
1 parent 8235d0c commit e5979a4
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 4 deletions.
6 changes: 4 additions & 2 deletions mne/cuda.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import numpy as np
from scipy.fft import irfft, rfft


from .utils import (
_check_option,
_explain_exception,
Expand All @@ -19,6 +18,7 @@

_cuda_capable = False


def get_shared_mem(
shape,
dtype=np.float64,
Expand All @@ -45,10 +45,11 @@ def get_shared_mem(
Returns
-------
a mapped array: np.ndarray
a mapped array: np.ndarray
An array to be passed into cupy.asarray, which does not copy if shared memory is already allocated.
"""
from numba import cuda

return cuda.mapped_array(
shape,
dtype=dtype,
Expand All @@ -59,6 +60,7 @@ def get_shared_mem(
wc=wc,
)


def get_cuda_memory(kind="available"):
"""Get the amount of free memory for CUDA operations.
Expand Down
4 changes: 2 additions & 2 deletions mne/tests/test_filter.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@
from scipy.signal import resample as sp_resample

from mne import Epochs, create_info
from mne.cuda import get_shared_mem
from mne._fiff.pick import _DATA_CH_TYPES_SPLIT
from mne.cuda import get_shared_mem
from mne.filter import (
_length_factors,
_overlap_add_filter,
Expand Down Expand Up @@ -858,7 +858,7 @@ def test_cuda_resampling():
for fro, to in ((1, 2), (2, 1), (1, 3), (3, 1)):
a1 = resample(a, fro, to, n_jobs=None, npad="auto", window=window)
x = get_shared_mem(a.shape)
x[:] = a
x[:] = a
a2 = resample(a, fro, to, n_jobs="cuda", npad="auto", window=window)
assert_allclose(a1, a2, rtol=1e-7, atol=1e-14)
assert_array_almost_equal(a1, a2, 14)
Expand Down

0 comments on commit e5979a4

Please sign in to comment.