Repository: ivy-llc/ivy

Issue:
embedding



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
@@ -1,5 +1,5 @@
 """Collection of Ivy neural network layers as stateful classes."""
-
+# flake8: noqa
 # local
 import ivy
 from ivy.func_wrapper import handle_nestable
@@ -1935,9 +1935,9 @@ def _forward(self, x):
 class FFT(Module):
     def __init__(
         self,
+        dim,
         /,
         *,
-        dim=0,
         norm="backward",
         n=None,
         out=None,
@@ -2095,3 +2095,43 @@ def _forward(self, x):
             axis=self.axis,
             norm=self.norm,
         )
+
+
+class Embedding(Module):
+    def __init__(self, indices, /, *, max_norm=None, out=None, device=None, dtype=None):
+        """
+        Class for embedding indices into a dense representation.
+
+        Parameters
+        ----------
+        indices
+            The indices to embed.
+        max_norm
+            If given, each embedding vector with norm larger than max_norm is renormalized to have norm max_norm.
+        out
+            If given, the result will be inserted into this tensor. Default: None
+        """
+        self._indices = indices
+        self._max_norm = max_norm
+        self._out = out
+        Module.__init__(self, device=device, dtype=dtype)
+
+    def _forward(self, inputs):
+        """
+        Forward pass of the layer.
+
+        Parameters
+        ----------
+        inputs
+            The input array to the layer.
+
+        Returns
+        -------
+            The output array of the layer.
+        """
+        return ivy.embedding(
+            inputs,
+            self._indices,
+            max_norm=self._max_norm,
+            out=self._out,
+        )


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
@@ -1493,3 +1493,46 @@ def test_dct(
         test_gradients=test_gradients,
         on_device=on_device,
     )
+
+
+# Embedding
+@handle_method(
+    method_tree="Embedding.__call__",
+    dtypes_indices_weights=helpers.embedding_helper(),
+    max_norm=st.one_of(st.none(), st.floats(min_value=1, max_value=5)),
+    number_positional_args=st.just(2),
+)
+def test_embedding_layer(
+    *,
+    dtypes_indices_weights,
+    max_norm,
+    number_positional_args,
+    test_gradients,
+    on_device,
+    class_name,
+    method_name,
+    ground_truth_backend,
+    init_flags,
+    method_flags,
+):
+    dtypes, indices, weights = dtypes_indices_weights
+    dtypes = [dtypes[1], dtypes[0]]
+
+    helpers.test_method(
+        ground_truth_backend=ground_truth_backend,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        init_all_as_kwargs_np={
+            "indices": indices,
+            "max_norm": max_norm,
+            "device": on_device,
+            "dtype": dtypes[0],
+        },
+        method_input_dtypes=dtypes,
+        method_all_as_kwargs_np={"inputs": weights},
+        class_name=class_name,
+        method_name=method_name,
+        test_gradients=test_gradients,
+        on_device=on_device,
+        number_positional_args=number_positional_args,
+    )