Skip to content

Commit 943b0a7

Browse files
committed
[BugFix] Robustize cuFINUFFT Python guru interface.
As per the docs: * finufft.Plan(n_modes_or_dim) accepts (int, tuple[int]) * cufinufft.Plan(n_modes) accepts (int, tuple[int]) In practice `n_modes_or_dim` is transformed internally to a NumPy array, hence providing any iterable to finufft.Plan() works. However `n_modes` in cufinufft.Plan() is more restrictive: it works with tuple[int], fails planning if list[int], and wrongfully passes planning if ndarray[int], but then seg-faults at execute(). This commit increases the robustness of cufinufft.Plan() to match the flexibility of the CPU interface.
1 parent 9744d14 commit 943b0a7

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

python/cufinufft/cufinufft/_plan.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
"""
66

77
import atexit
8+
import collections.abc
9+
import numbers
810
import sys
911
import warnings
1012

@@ -108,8 +110,12 @@ def __init__(self, nufft_type, n_modes, n_trans=1, eps=1e-6, isign=None,
108110
else:
109111
raise TypeError("Expected complex64 or complex128.")
110112

111-
if isinstance(n_modes, int):
113+
if isinstance(n_modes, numbers.Integral):
112114
n_modes = (n_modes,)
115+
elif isinstance(n_modes, collections.abc.Iterable):
116+
n_modes = tuple(n_modes)
117+
else:
118+
raise ValueError(f"Invalid n_modes '{n_modes}'")
113119

114120
self.dim = len(n_modes)
115121
self.type = nufft_type

0 commit comments

Comments
 (0)