FFT: finish up the dir(np.fft) contents

Author: ev-br

torch_np/fft.py

@@ -106,3 +106,23 @@ def hfft(a: ArrayLike, n=None, axis=-1, norm=None):
 @upcast
 def ihfft(a: ArrayLike, n=None, axis=-1, norm=None):
     return torch.fft.ihfft(a, n, dim=axis, norm=norm)
+
+
+@normalizer
+def fftfreq(n, d=1.0):
+    return torch.fft.fftfreq(n, d)
+
+
+@normalizer
+def rfftfreq(n, d=1.0):
+    return torch.fft.rfftfreq(n, d)
+
+
+@normalizer
+def fftshift(x: ArrayLike, axes=None):
+    return torch.fft.fftshift(x, axes)
+
+
+@normalizer
+def ifftshift(x: ArrayLike, axes=None):
+    return torch.fft.ifftshift(x, axes)

torch_np/tests/numpy_tests/fft/test_helper.py

@@ -85,7 +85,7 @@ def test_uneven_dims(self):
 
     def test_equal_to_original(self):
         """ Test that the new (>=v1.15) implementation (see #10073) is equal to the original (<=v1.14) """
-        from numpy.core import asarray, concatenate, arange, take
+        from torch_np import asarray, concatenate, arange, take
 
         def original_fftshift(x, axes=None):
             """ How fftshift was implemented in v1.14"""