BUG: Fix fftn failure when both out and s parameters are given #28895
Conversation
|
Looks like a pypy issue. Similar to #28469 what they encountered. |
numpy/fft/_pocketfft.py
Outdated
| if out is not None and s is not None and any(s[i] != a.shape[axes[i]] for i in range(len(axes))): | ||
| intermediate_out = None | ||
| else: | ||
| intermediate_out = out |
There was a problem hiding this comment.
Does it bring any performance improvement if we pass a itself when possible as out in intermediate steps (when out is originally None)? We might not be able to do so in the first iteration tough.
mhvk
left a comment
There was a problem hiding this comment.
@shibozhou - I had some comments and then realized it may be possible to be a bit smarter about intermediate results, and at the same time address @vtavana's comment about using in-place with the input where possible (which should also speed up the out=None case). I'll push to your branch with that, but the tests still need to be improved a bit. Furthermore, the documentation for out is now no longer correct, so that needs adjusting as well.
numpy/fft/_pocketfft.py
Outdated
| def _raw_fftnd(a, s=None, axes=None, function=fft, norm=None, out=None): | ||
| a = asarray(a) | ||
| s, axes = _cook_nd_args(a, s, axes) | ||
| s, axes = _cook_nd_args(a, s, axes, function is irfft) |
There was a problem hiding this comment.
Why this change? _raw_fftnd is never called with function=irfft.
numpy/fft/_pocketfft.py
Outdated
| s, axes = _cook_nd_args(a, s, axes) | ||
| s, axes = _cook_nd_args(a, s, axes, function is irfft) | ||
|
|
||
| if out is not None and s is not None and any(s[i] != a.shape[axes[i]] for i in range(len(axes))): |
There was a problem hiding this comment.
Note that s is already overridden here.
numpy/fft/tests/test_pocketfft.py
Outdated
| assert result is out | ||
| assert_array_equal(result, expected) | ||
|
|
||
| def test_fftn_with_out_and_s(self): |
There was a problem hiding this comment.
I'd suggest to use a pytest.mark.parametrize here to do the two tests and add more possible s tuples, making values both bigger and smaller than the array in various positions. Also, could you move the test right below test_fftn_out_and_s_interaction (and maybe call it test_fftn_out_and_s_interaction2).
numpy/fft/tests/test_pocketfft.py
Outdated
| def test_fftn_with_out_and_s(self): | ||
| # fftn works correctly with both out and s parameters | ||
| rng = np.random.RandomState(42) | ||
| x = np.random.rand(4, 4, 4) + 1j * np.random.rand(4, 4, 4) |
There was a problem hiding this comment.
It may be wise to ensure we never modify the input, by doing after the definition, x.flags.writeable = False.
numpy/fft/_pocketfft.py
Outdated
| elif this_out is not None and n < this_out.shape[axis]: | ||
| # For an intermediate step where we return fewer elements, we | ||
| # can use a smaller view of the previous array. | ||
| this_out = this_out[(slice(None,),)*(axis-1) + (slice(n),)] |
There was a problem hiding this comment.
I believe it should be axis
| this_out = this_out[(slice(None,),)*(axis-1) + (slice(n),)] | |
| this_out = this_out[(slice(None,),)*axis + (slice(n),)] |
There was a problem hiding this comment.
Yes, you're right - which just goes to show that this really needs additional tests!
@shibozhou - are you up for adding tests, or would you prefer that I do it?
Along a similar vein: I think axis can actually be negative here -- they should be normalized, perhaps best in _cook_nd_args.
mhvk
left a comment
There was a problem hiding this comment.
@shibozhou - sorry for the slow response. The code looks mostly good -- I have just one comment -- but the tests fail so obviously something is wrong. I didn't have time to look in detail, but it seems related to the array shape. Probably easiest to investigate locally, in debug mode.
| itl.reverse() | ||
| for ii in itl: | ||
| a = function(a, n=s[ii], axis=axes[ii], norm=norm, out=out) | ||
| axes = [ax % a.ndim for ax in axes] |
There was a problem hiding this comment.
Please use normalize_axis_index(ax, a.ndim) here - it is already imported on top (it ensures that ax > a.ndim or ax < -a.ndim properly error).
Fixes #28890
Problem
When N-dim FFTs are calculated using
fftn, the code performs a loop over all axes, applying 1D FFts in sequence. However, when bothoutandsare specified, intermediate transforms change the array shape, making theoutparameter incompatible with these intermediate results.Solution
This pr changes the
_raw_fftndto only use theoutparameter for the final transformation, not for intermediate steps. The current function works correctly when bothoutandshave different shapes than the input array.Test
The added test case verifies the fix with an input array of shape
(4, 4, 4)and an output with shape(6, 6, 6)