fft2D' and fft3D' fail on arrays where dimension sizes are different

[sh54]

In fft2D'

        fft' a  = A.transpose . fft sign (Z:.width)  height
              >-> A.transpose . fft sign (Z:.height) width
                $ a

should be

        fft' a  = A.transpose . fft sign (Z:.height) width
              >-> A.transpose . fft sign (Z:.width)  height
                $ a

Otherwise an index error is given if supplied some array with a shape like Z :. 8 :. 4.

In fft3D'

scale = P.fromIntegral (width * height) should be scale = P.fromIntegral (depth * width * height).

Also:

        fft' a  = rotate3D . fft sign (Z:.width :.depth)  height
              >-> rotate3D . fft sign (Z:.height:.width)  depth
              >-> rotate3D . fft sign (Z:.depth :.height) width
                $ a

should be

        fft' a  = rotate3D . fft sign (Z:.depth :.height) width
              >-> rotate3D . fft sign (Z:.height:.width)  depth
              >-> rotate3D . fft sign (Z:.width :.depth)  height
                $ a

Otherwise an index error is given if supplied some array with a shape like Z :. 8 :. 4 :. 2.

Also there seems to be an issue with the backpermute call in rotate3D when given an array with a shape like Z :. 8 :. 4 :. 2 so I have used the following replacement as a work around:

rotate3D :: Elt e => Acc (Array DIM3 e) -> Acc (Array DIM3 e)
rotate3D arr = do
  let Z :. k :. l :. m = unlift $ A.shape arr  :: Z :. Exp Int :. Exp Int :. Exp Int
  let sz = lift $ Z :. l :. m :. k
  generate sz $ \ix -> do
    let Z :. l' :. m' :. k' = unlift ix :: Z :. Exp Int :. Exp Int :. Exp Int
    arr ! (index3 k' l' m')

The backpermute issue has been registered here.

I checked my changes against the equivalent numpy code and got the same results. I also quickly checked doing an inverse fft on the fft and got back the original arrays.

This just is an issue for the non-accelerated versions.

[tmcdonell]

Thanks! I've added both of these fixes.