Repository: ivy-llc/ivy

Issue:
fft



Discussion:

Resolving patch:
diff --git a/ivy/stateful/layers.py b/ivy/stateful/layers.py
--- a/ivy/stateful/layers.py
+++ b/ivy/stateful/layers.py
@@ -2,10 +2,9 @@
 
 # local
 import ivy
-from ivy.stateful.module import Module
-from ivy.stateful.initializers import Zeros, GlorotUniform
 from ivy.func_wrapper import handle_nestable
-
+from ivy.stateful.initializers import GlorotUniform, Zeros
+from ivy.stateful.module import Module
 
 # ToDo: update docstrings and typehints according to ivy\layers
 
@@ -1933,6 +1932,61 @@ def _forward(self, x):
         )
 
 
+class FFT(Module):
+    def __init__(
+        self,
+        /,
+        *,
+        dim=0,
+        norm="backward",
+        n=None,
+        out=None,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Class for applying FFT to input.
+
+        Parameters
+        ----------
+        dim : int
+            Dimension along which to take the FFT.
+        norm : str
+            Normalization mode. Default: 'backward'
+        n : int
+            Size of the FFT. Default: None
+        out : int
+            Size of the output. Default: None
+        """
+        self._dim = dim
+        self._norm = norm
+        self._n = n
+        self._out = out
+        Module.__init__(self, device=device, dtype=dtype)
+
+    def _forward(self, inputs):
+        """
+        Forward pass of the layer.
+
+        Parameters
+        ----------
+        inputs : array
+            Input array to take the FFT of.
+
+        Returns
+        -------
+        array
+            The output array of the layer.
+        """
+        return ivy.fft(
+            inputs,
+            self._dim,
+            norm=self._norm,
+            n=self._n,
+            out=self._out,
+        )
+
+
 class AvgPool1D(Module):
     def __init__(
         self,


Test patch:
diff --git a/ivy_tests/test_ivy/test_stateful/test_layers.py b/ivy_tests/test_ivy/test_stateful/test_layers.py
--- a/ivy_tests/test_ivy/test_stateful/test_layers.py
+++ b/ivy_tests/test_ivy/test_stateful/test_layers.py
@@ -2,16 +2,20 @@
 
 # global
 import numpy as np
-from hypothesis import strategies as st, assume
+from hypothesis import assume
+from hypothesis import strategies as st
 
 # local
 import ivy
-from ivy.functional.ivy.layers import _deconv_length
-from ivy.functional.ivy.gradients import _variable
-from ivy.data_classes.container import Container
 import ivy_tests.test_ivy.helpers as helpers
-from ivy_tests.test_ivy.helpers.assertions import assert_same_type_and_shape
+from ivy.data_classes.container import Container
+from ivy.functional.ivy.gradients import _variable
+from ivy.functional.ivy.layers import _deconv_length
 from ivy_tests.test_ivy.helpers import handle_method
+from ivy_tests.test_ivy.helpers.assertions import assert_same_type_and_shape
+from ivy_tests.test_ivy.test_functional.test_experimental.test_nn import (
+    test_layers as exp_layers_tests,
+)
 from ivy_tests.test_ivy.test_functional.test_experimental.test_nn.test_layers import (
     valid_dct,
 )
@@ -1382,6 +1386,43 @@ def test_adaptive_avg_pool1d_layer(
     )
 
 
+# FFT
+@handle_method(
+    method_tree="FFT.__call__",
+    x_and_fft=exp_layers_tests.x_and_fft(),
+)
+def test_fft_layer(
+    *,
+    x_and_fft,
+    test_gradients,
+    on_device,
+    class_name,
+    method_name,
+    ground_truth_backend,
+    init_flags,
+    method_flags,
+):
+    dtype, x, dim, norm, n = x_and_fft
+    helpers.test_method(
+        ground_truth_backend=ground_truth_backend,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        init_all_as_kwargs_np={
+            "dim": dim,
+            "norm": norm,
+            "n": n,
+            "device": on_device,
+            "dtype": dtype[0],
+        },
+        method_input_dtypes=dtype,
+        method_all_as_kwargs_np={"inputs": x[0]},
+        class_name=class_name,
+        method_name=method_name,
+        test_gradients=test_gradients,
+        on_device=on_device,
+    )
+
+
 # AvgPool1D
 @handle_method(
     method_tree="AvgPool1D.__call__",