Repository: ivy-llc/ivy

Issue:
dct



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
@@ -1983,3 +1983,61 @@ def _forward(self, inputs):
             self._padding,
             data_format=self._data_format,
         )
+
+
+class Dct(Module):
+    def __init__(
+        self,
+        *,
+        type=2,
+        n=None,
+        axis=-1,
+        norm=None,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Class for applying the Discrete Cosine Transform over mini-batch of inputs.
+
+        Parameters
+        ----------
+        x
+            The input signal.
+        type
+            The type of the dct. Must be 1, 2, 3 or 4.
+        n
+            The length of the transform. If n is less than the input signal lenght,
+            then x is truncated, if n is larger then x is zero-padded.
+        axis
+            The axis to compute the DCT along.
+        norm
+            The type of normalization to be applied. Must be either None or "ortho".
+        device
+            device on which to create the layer's variables 'cuda:0', 'cuda:1', 'cpu'
+        """
+        self.type = type
+        self.n = n
+        self.axis = axis
+        self.norm = norm
+        Module.__init__(self, device=device, dtype=dtype)
+
+    def _forward(self, x):
+        """
+        Forward pass of the layer.
+
+        Parameters
+        ----------
+        x
+            The input array to the layer.
+
+        Returns
+        -------
+            The output array of the layer.
+        """
+        return ivy.dct(
+            x,
+            type=self.type,
+            n=self.n,
+            axis=self.axis,
+            norm=self.norm,
+        )


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
@@ -12,7 +12,9 @@
 import ivy_tests.test_ivy.helpers as helpers
 from ivy_tests.test_ivy.helpers.assertions import assert_same_type_and_shape
 from ivy_tests.test_ivy.helpers import handle_method
-
+from ivy_tests.test_ivy.test_functional.test_experimental.test_nn.test_layers import (
+    valid_dct,
+)
 
 # Helpers #
 # --------#
@@ -1413,3 +1415,40 @@ def test_avgpool1d_layer(
         test_gradients=test_gradients,
         on_device=on_device,
     )
+
+
+@handle_method(
+    method_tree="Dct.__call__",
+    dtype_x_and_args=valid_dct(),
+)
+def test_dct(
+    *,
+    dtype_x_and_args,
+    test_gradients,
+    on_device,
+    class_name,
+    method_name,
+    ground_truth_backend,
+    init_flags,
+    method_flags,
+):
+    dtype, x, type, n, axis, norm = dtype_x_and_args
+    helpers.test_method(
+        ground_truth_backend=ground_truth_backend,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        init_all_as_kwargs_np={
+            "dtype": dtype[0],
+            "type": type,
+            "n": n,
+            "axis": axis,
+            "norm": norm,
+            "device": on_device,
+        },
+        method_input_dtypes=dtype,
+        method_all_as_kwargs_np={"x": x[0]},
+        class_name=class_name,
+        method_name=method_name,
+        test_gradients=test_gradients,
+        on_device=on_device,
+    )