mkl_fft._pydfti.rfftn_numpy error when using axes specification

[Jeitan]

I originally posted this as a numpy issue, but after some research it appears to be something wrong with mkl_fft, although it only happens under Python 3 (and only in later releases, where "later" means past whatever the default was in Anaconda 5.0.1).

You can see some background on the issue here - the context is that I use numpy's rfftn to make a PSD of simulated ocean waves.

Long story short, when I make a conda environment like this:

> conda create --name test1 anaconda=2019.03 python=2

and this:

> conda create --name test2 anaconda=2019.03 python=3

If I have a numpy array called 'waves' that is dimensioned [64,512,512] and try to run the following:

import numpy.fft as nfft
realgood2 = nfft.rfftn(waves, axes=[2,1,0])

In test environment 1 it works perfectly. In test environment 2, I get this:

Traceback (most recent call last):
  File "gist.py", line 57, in <module>
    realgood2 = nfft.rfftn(waves, axes=[2,1,0])
  File "C:\Users\asmith\AppData\Local\Continuum\anaconda3\envs\3_2019.03\lib\site-packages\mkl_fft\_numpy_fft.py", line 1043, in rfftn
    output = mkl_fft.rfftn_numpy(a, s, axes)
  File "mkl_fft\_pydfti.pyx", line 951, in mkl_fft._pydfti.rfftn_numpy
ValueError: could not broadcast input array from shape (512,64) into shape (512,512)

[Jeitan]

After some investigation, I actually have no idea how this changes between py2 and py3, or actually why it ever works, because I may have found a logic bug. Someone should definitely verify (or disprove), because I didn't even know numpy used mkl or mkl_fft under the hood until today, so I am very unfamiliar with the code. I am, however, familiar with FFTs, and axes specification.

In _pydfti.pyx line 941, new arrays are constructed for "s" and "axes" called "ss" and "aa", which are designed to be passed to the fftn call for the remaining axes (the 1D rfft was called above, on the correct axis). s and axes are passed into something called _remove_axis which looks like it's designed to remove the axis that we've already FFTed over so as to do the full complex fft on the other axes (which happens in a loop over the rfft axis, which makes perfect sense).

At this point, both s and axes are, by construction, defined to have the rfft axis be the last axis (axes[-1] and s[-1] do match). However, the axis to be removed that is passed in is the variable la, which is the specification of the rfft axis (in my case above, it is 0). But if I'm reading how it works right, what should get passed right there should always be -1.

Just looking at the code, it's actually sort of a mystery why it doesn't throw the same error when I run it under Python 2 - it does throw something ("all axes must be unique") when I dive down and run _fftnd_impl by hand. But yeah, if we can get that la changed to -1, I think that is correct?