adding ifft2 method to ops (#20447)

* adding ifft2 method to ops

* fixes all test checks

* using built-in versions in backends

Author: rohithpudari

keras/api/_tf_keras/keras/ops/__init__.py

@@ -47,6 +47,7 @@
 from keras.src.ops.math import extract_sequences
 from keras.src.ops.math import fft
 from keras.src.ops.math import fft2
+from keras.src.ops.math import ifft2
 from keras.src.ops.math import in_top_k
 from keras.src.ops.math import irfft
 from keras.src.ops.math import istft

keras/api/ops/__init__.py

@@ -47,6 +47,7 @@
 from keras.src.ops.math import extract_sequences
 from keras.src.ops.math import fft
 from keras.src.ops.math import fft2
+from keras.src.ops.math import ifft2
 from keras.src.ops.math import in_top_k
 from keras.src.ops.math import irfft
 from keras.src.ops.math import istft

keras/src/backend/jax/math.py

@@ -123,6 +123,12 @@ def fft2(x):
     return jnp.real(complex_output), jnp.imag(complex_output)
 
 
+def ifft2(x):
+    complex_input = _get_complex_tensor_from_tuple(x)
+    complex_output = jnp.fft.ifft2(complex_input)
+    return jnp.real(complex_output), jnp.imag(complex_output)
+
+
 def rfft(x, fft_length=None):
     complex_output = jnp.fft.rfft(x, n=fft_length, axis=-1, norm="backward")
     return jnp.real(complex_output), jnp.imag(complex_output)

keras/src/backend/numpy/math.py

@@ -144,6 +144,12 @@ def fft2(x):
     return np.array(real), np.array(imag)
 
 
+def ifft2(x):
+    complex_input = _get_complex_tensor_from_tuple(x)
+    complex_output = np.fft.ifft2(complex_input)
+    return np.real(complex_output), np.imag(complex_output)
+
+
 def rfft(x, fft_length=None):
     complex_output = np.fft.rfft(x, n=fft_length, axis=-1, norm="backward")
     # numpy always outputs complex128, so we need to recast the dtype

keras/src/backend/tensorflow/math.py

@@ -113,6 +113,15 @@ def fft2(x):
     return tf.math.real(complex_output), tf.math.imag(complex_output)
 
 
+def ifft2(x):
+    real, imag = x
+    h = cast(tf.shape(real)[-2], "float32")
+    w = cast(tf.shape(real)[-1], "float32")
+    real_conj, imag_conj = real, -imag
+    fft_real, fft_imag = fft2((real_conj, imag_conj))
+    return fft_real / (h * w), -fft_imag / (h * w)
+
+
 def rfft(x, fft_length=None):
     if fft_length is not None:
         fft_length = [fft_length]

keras/src/backend/torch/math.py

@@ -203,6 +203,12 @@ def fft2(x):
     return complex_output.real, complex_output.imag
 
 
+def ifft2(x):
+    complex_input = _get_complex_tensor_from_tuple(x)
+    complex_output = torch.fft.ifft2(complex_input)
+    return complex_output.real, complex_output.imag
+
+
 def rfft(x, fft_length=None):
     x = convert_to_tensor(x)
     complex_output = torch.fft.rfft(x, n=fft_length, dim=-1, norm="backward")

keras/src/ops/math.py

@@ -473,6 +473,81 @@ def fft2(x):
     return backend.math.fft2(x)
 
 
+class IFFT2(Operation):
+    def __init__(self):
+        super().__init__()
+        self.axes = (-2, -1)
+
+    def compute_output_spec(self, x):
+        if not isinstance(x, (tuple, list)) or len(x) != 2:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                f"imaginary. Received: x={x}"
+            )
+
+        real, imag = x
+        # Both real and imaginary parts should have the same shape.
+        if real.shape != imag.shape:
+            raise ValueError(
+                "Input `x` should be a tuple of two tensors - real and "
+                "imaginary. Both the real and imaginary parts should have the "
+                f"same shape. Received: x[0].shape = {real.shape}, "
+                f"x[1].shape = {imag.shape}"
+            )
+        # We are calculating 2D IFFT. Hence, rank >= 2.
+        if len(real.shape) < 2:
+            raise ValueError(
+                f"Input should have rank >= 2. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        # The axes along which we are calculating IFFT should be fully-defined.
+        m = real.shape[self.axes[0]]
+        n = real.shape[self.axes[1]]
+        if m is None or n is None:
+            raise ValueError(
+                f"Input should have its {self.axes} axes fully-defined. "
+                f"Received: input.shape = {real.shape}"
+            )
+
+        return (
+            KerasTensor(shape=real.shape, dtype=real.dtype),
+            KerasTensor(shape=imag.shape, dtype=imag.dtype),
+        )
+
+    def call(self, x):
+        return backend.math.ifft2(x)
+
+
+@keras_export("keras.ops.ifft2")
+def ifft2(x):
+    """Computes the 2D Inverse Fast Fourier Transform along the last two axes of
+        input.
+
+    Args:
+        x: Tuple of the real and imaginary parts of the input tensor. Both
+            tensors in the tuple should be of floating type.
+
+    Returns:
+        A tuple containing two tensors - the real and imaginary parts of the
+        output.
+
+    Example:
+
+    >>> x = (
+    ...     keras.ops.convert_to_tensor([[1., 2.], [2., 1.]]),
+    ...     keras.ops.convert_to_tensor([[0., 1.], [1., 0.]]),
+    ... )
+    >>> ifft2(x)
+    (array([[ 6.,  0.],
+        [ 0., -2.]], dtype=float32), array([[ 2.,  0.],
+        [ 0., -2.]], dtype=float32))
+    """
+    if any_symbolic_tensors(x):
+        return IFFT2().symbolic_call(x)
+    return backend.math.ifft2(x)
+
+
 class RFFT(Operation):
     def __init__(self, fft_length=None):
         super().__init__()

keras/src/ops/math_test.py

@@ -218,6 +218,15 @@ def test_fft2(self):
         self.assertEqual(real_output.shape, ref_shape)
         self.assertEqual(imag_output.shape, ref_shape)
 
+    def test_ifft2(self):
+        real = KerasTensor((None, 4, 3), dtype="float32")
+        imag = KerasTensor((None, 4, 3), dtype="float32")
+        real_output, imag_output = kmath.ifft2((real, imag))
+        ref = np.fft.ifft2(np.ones((2, 4, 3)))
+        ref_shape = (None,) + ref.shape[1:]
+        self.assertEqual(real_output.shape, ref_shape)
+        self.assertEqual(imag_output.shape, ref_shape)
+
     @parameterized.parameters([(None,), (1,), (5,)])
     def test_rfft(self, fft_length):
         x = KerasTensor((None, 4, 3), dtype="float32")
@@ -354,6 +363,14 @@ def test_fft2(self):
         self.assertEqual(real_output.shape, ref.shape)
         self.assertEqual(imag_output.shape, ref.shape)
 
+    def test_ifft2(self):
+        real = KerasTensor((2, 4, 3), dtype="float32")
+        imag = KerasTensor((2, 4, 3), dtype="float32")
+        real_output, imag_output = kmath.ifft2((real, imag))
+        ref = np.fft.ifft2(np.ones((2, 4, 3)))
+        self.assertEqual(real_output.shape, ref.shape)
+        self.assertEqual(imag_output.shape, ref.shape)
+
     def test_rfft(self):
         x = KerasTensor((2, 4, 3), dtype="float32")
         real_output, imag_output = kmath.rfft(x)
@@ -715,6 +732,18 @@ def test_fft2(self):
         self.assertAllClose(real_ref, real_output)
         self.assertAllClose(imag_ref, imag_output)
 
+    def test_ifft2(self):
+        real = np.random.random((2, 4, 3)).astype(np.float32)
+        imag = np.random.random((2, 4, 3)).astype(np.float32)
+        complex_arr = real + 1j * imag
+
+        real_output, imag_output = kmath.ifft2((real, imag))
+        ref = np.fft.ifft2(complex_arr)
+        real_ref = np.real(ref)
+        imag_ref = np.imag(ref)
+        self.assertAllClose(real_ref, real_output)
+        self.assertAllClose(imag_ref, imag_output)
+
     @parameterized.parameters([(None,), (3,), (15,)])
     def test_rfft(self, n):
         # Test 1D.