Repository: google/jax
Issue: SIGFPE + RuntimeError caused by missing shape rule for FFT

The following code

```python
from jax import lax 
from jaxlib import xla_client
import numpy as np

x = np.ones((14, 15))
kwargs = { "fft_type": xla_client.FftType.FFT
         , "fft_lengths": [0]
         }
         
lax.fft(x, **kwargs)
```

triggers a floating point exception and thus crashes.

Changed files:
jax/_src/lax/fft.py, tests/fft_test.py

diff --git a/jax/_src/lax/fft.py b/jax/_src/lax/fft.py
index 96230fa8f..8e3490873 100644
--- a/jax/_src/lax/fft.py
+++ b/jax/_src/lax/fft.py
@@ -39,13 +39,13 @@ __all__ = [
 ]
 
 def _str_to_fft_type(s: str) -> xla_client.FftType:
-  if s == "FFT":
+  if s in ("fft", "FFT"):
     return xla_client.FftType.FFT
-  elif s == "IFFT":
+  elif s in ("ifft", "IFFT"):
     return xla_client.FftType.IFFT
-  elif s == "RFFT":
+  elif s in ("rfft", "RFFT"):
     return xla_client.FftType.RFFT
-  elif s == "IRFFT":
+  elif s in ("irfft", "IRFFT"):
     return xla_client.FftType.IRFFT
   else:
     raise ValueError(f"Unknown FFT type '{s}'")
@@ -79,14 +79,27 @@ _real_dtype = lambda dtype: np.finfo(dtype).dtype
 _is_even = lambda x: x % 2 == 0
 
 def fft_abstract_eval(x, fft_type, fft_lengths):
+  if len(fft_lengths) > x.ndim:
+    raise ValueError(f"FFT input shape {x.shape} must have at least as many "
+                    f"input dimensions as fft_lengths {fft_lengths}.")
   if fft_type == xla_client.FftType.RFFT:
+    if x.shape[-len(fft_lengths):] != fft_lengths:
+      raise ValueError(f"RFFT input shape {x.shape} minor dimensions must "
+                      f"be equal to fft_lengths {fft_lengths}")
     shape = (x.shape[:-len(fft_lengths)] + fft_lengths[:-1]
              + (fft_lengths[-1] // 2 + 1,))
     dtype = _complex_dtype(x.dtype)
   elif fft_type == xla_client.FftType.IRFFT:
+    if x.shape[-len(fft_lengths):-1] != fft_lengths[:-1]:
+      raise ValueError(f"IRFFT input shape {x.shape} minor dimensions must "
+                      "be equal to all except the last fft_length, got "
+                      f"fft_lengths={fft_lengths}")
     shape = x.shape[:-len(fft_lengths)] + fft_lengths
     dtype = _real_dtype(x.dtype)
   else:
+    if x.shape[-len(fft_lengths):] != fft_lengths:
+      raise ValueError(f"FFT input shape {x.shape} minor dimensions must "
+                      f"be equal to fft_lengths {fft_lengths}")
     shape = x.shape
     dtype = x.dtype
   return x.update(shape=shape, dtype=dtype)
diff --git a/tests/fft_test.py b/tests/fft_test.py
index 590a64b5c..190ef70bb 100644
--- a/tests/fft_test.py
+++ b/tests/fft_test.py
@@ -104,6 +104,30 @@ class FftTest(jtu.JaxTestCase):
     x = rng((10,), np.complex64)
     self.assertAllClose(np.fft.fft(x).astype(np.complex64),
                         lax.fft(x, "FFT", fft_lengths=(10,)))
+    self.assertAllClose(np.fft.fft(x).astype(np.complex64),
+                        lax.fft(x, "fft", fft_lengths=(10,)))
+
+  def testLaxFftErrors(self):
+    with self.assertRaises(
+      ValueError,
+      msg="FFT input shape (14, 15) must have at least as many input "
+          "dimensions as fft_lengths (4, 5, 6)"):
+      lax.fft(np.ones((14, 15)), fft_type="fft", fft_lengths=(4, 5, 6))
+    with self.assertRaises(
+      ValueError,
+      msg="FFT input shape (14, 15) minor dimensions must be equal to "
+          "fft_lengths (17,)"):
+      lax.fft(np.ones((14, 15)), fft_type="fft", fft_lengths=(17,))
+    with self.assertRaises(
+      ValueError,
+      msg="RFFT input shape (14, 15) minor dimensions must be equal to "
+          "fft_lengths (14, 15,)"):
+      lax.fft(np.ones((2, 14, 15)), fft_type="rfft", fft_lengths=(14, 12))
+    with self.assertRaises(
+      ValueError,
+      msg="IRFFT input shape (14, 15) minor dimensions must be equal to "
+          "all except the last fft_length, got fft_lengths=(14, 15,)"):
+      lax.fft(np.ones((2, 14, 15)), fft_type="irfft", fft_lengths=(13, 15))
 
   @parameterized.parameters((np.float32,), (np.float64,))
   def testLaxIrfftDoesNotMutateInputs(self, dtype):