Add axes= keyword to fft dtype sample computation

This is important for irfft2 dtype analysis
mrocklin · Aug 21, 2018 · 5b25a09 · 5b25a09
1 parent 9636454
commit 5b25a09
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/dask/array/fft.py b/dask/array/fft.py
@@ -169,7 +169,10 @@ def func(a, s=None, axes=None):
         _dtype = dtype
         if _dtype is None:
             sample = np.ones(a.ndim * (8,), dtype=a.dtype)
-            _dtype = fft_func(sample).dtype
+            try:
+                _dtype = fft_func(sample, axes=axes).dtype
+            except TypeError:
+                _dtype = fft_func(sample).dtype
 
         for each_axis in axes:
             if len(a.chunks[each_axis]) != 1: