Add embedding tests

backends/test/compliance_suite/__init__.py

@@ -114,13 +114,14 @@ def wrapped_test(self):
 
 
 class OperatorTest(unittest.TestCase):
-    def _test_op(self, model, inputs, tester_factory):
+    def _test_op(self, model, inputs, tester_factory, use_random_test_inputs=True):
         tester = (
             tester_factory(
                 model,
                 inputs,
             )
             .export()
+            .dump_artifact()
             .to_edge_transform_and_lower()
         )
 
@@ -136,6 +137,11 @@ def _test_op(self, model, inputs, tester_factory):
                 tester
                 .to_executorch()
                 .serialize()
-                .run_method_and_compare_outputs()
+                # If use_random_test_inputs is False, explicitly pass the export inputs for the test.
+                # This is useful for ops like embedding where the random input generation isn't aware
+                # of the constraints on the input data values (e.g. the indices must be within bounds).
+                # If use_random_test_inputs is True, a value of None is passed to cause test inputs to
+                # be randomly generated.
+                .run_method_and_compare_outputs(inputs = inputs if not use_random_test_inputs else None)
             )
 

backends/test/compliance_suite/operators/test_embedding.py

@@ -0,0 +1,68 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+from typing import Callable, Optional
+
+import torch
+
+from executorch.backends.test.compliance_suite import (
+    dtype_test,
+    operator_test,
+    OperatorTest,
+)
+
+class Model(torch.nn.Module):
+    def __init__(
+        self,
+        num_embeddings=10,
+        embedding_dim=5,
+        padding_idx: Optional[int] = None,
+        norm_type: float = 2.0,
+    ):
+        super().__init__()
+        self.embedding = torch.nn.Embedding(
+            num_embeddings=num_embeddings,
+            embedding_dim=embedding_dim,
+            padding_idx=padding_idx,
+            norm_type=norm_type,
+        )
+
+    def forward(self, x):
+        return self.embedding(x)
+
+@operator_test
+class TestEmbedding(OperatorTest):
+    @dtype_test
+    def test_embedding_dtype(self, dtype, tester_factory: Callable) -> None:
+        # Input shape: (batch_size, seq_length)
+        # Note: Input indices should be of type Long (int64)
+        model = Model().to(dtype)
+        self._test_op(model, (torch.randint(0, 10, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_basic(self, tester_factory: Callable) -> None:
+        # Basic test with default parameters
+        self._test_op(Model(), (torch.randint(0, 10, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_sizes(self, tester_factory: Callable) -> None:
+        # Test with different dictionary sizes and embedding dimensions
+        self._test_op(Model(num_embeddings=5, embedding_dim=3), 
+                     (torch.randint(0, 5, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+        self._test_op(Model(num_embeddings=100, embedding_dim=10), 
+                     (torch.randint(0, 100, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+        self._test_op(Model(num_embeddings=1000, embedding_dim=50), 
+                     (torch.randint(0, 1000, (2, 4), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_padding_idx(self, tester_factory: Callable) -> None:
+        # Test with padding_idx
+        self._test_op(Model(padding_idx=0), 
+                     (torch.randint(0, 10, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+        self._test_op(Model(padding_idx=5), 
+                     (torch.randint(0, 10, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_input_shapes(self, tester_factory: Callable) -> None:
+        # Test with different input shapes
+        self._test_op(Model(), (torch.randint(0, 10, (5,), dtype=torch.long),), tester_factory, use_random_test_inputs=False)  # 1D input
+        self._test_op(Model(), (torch.randint(0, 10, (2, 8), dtype=torch.long),), tester_factory, use_random_test_inputs=False)  # 2D input
+        self._test_op(Model(), (torch.randint(0, 10, (2, 3, 4), dtype=torch.long),), tester_factory, use_random_test_inputs=False)  # 3D input
+

backends/test/compliance_suite/operators/test_embedding_bag.py

@@ -0,0 +1,100 @@
+# (c) Meta Platforms, Inc. and affiliates. Confidential and proprietary.
+
+# pyre-strict
+
+from typing import Callable, Optional
+
+import torch
+
+from executorch.backends.test.compliance_suite import (
+    dtype_test,
+    operator_test,
+    OperatorTest,
+)
+
+class Model(torch.nn.Module):
+    def __init__(
+        self,
+        num_embeddings=10,
+        embedding_dim=5,
+        mode='mean',
+        padding_idx: Optional[int] = None,
+        norm_type: float = 2.0,
+        include_last_offset: bool = False,
+    ):
+        super().__init__()
+        self.embedding_bag = torch.nn.EmbeddingBag(
+            num_embeddings=num_embeddings,
+            embedding_dim=embedding_dim,
+            mode=mode,
+            padding_idx=padding_idx,
+            norm_type=norm_type,
+            include_last_offset=include_last_offset,
+        )
+
+    def forward(self, x, offsets=None):
+        return self.embedding_bag(x, offsets)
+
+@operator_test
+class TestEmbeddingBag(OperatorTest):
+    @dtype_test
+    def test_embedding_bag_dtype(self, dtype, tester_factory: Callable) -> None:
+        # Input: indices and offsets
+        # Note: Input indices should be of type Long (int64)
+        model = Model().to(dtype)
+        indices = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
+        offsets = torch.tensor([0, 4], dtype=torch.long)  # 2 bags
+        self._test_op(model, (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_bag_basic(self, tester_factory: Callable) -> None:
+        # Basic test with default parameters
+        indices = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
+        offsets = torch.tensor([0, 4], dtype=torch.long)  # 2 bags
+        self._test_op(Model(), (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_bag_sizes(self, tester_factory: Callable) -> None:
+        # Test with different dictionary sizes and embedding dimensions
+        indices = torch.tensor([1, 2, 3, 1], dtype=torch.long)
+        offsets = torch.tensor([0, 2], dtype=torch.long)
+
+        self._test_op(Model(num_embeddings=5, embedding_dim=3), 
+                     (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+        indices = torch.tensor([5, 20, 10, 43, 7], dtype=torch.long)
+        offsets = torch.tensor([0, 2, 4], dtype=torch.long)
+        self._test_op(Model(num_embeddings=50, embedding_dim=10), 
+                     (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+        indices = torch.tensor([100, 200, 300, 400], dtype=torch.long)
+        offsets = torch.tensor([0, 2], dtype=torch.long)
+        self._test_op(Model(num_embeddings=500, embedding_dim=20), 
+                     (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_bag_modes(self, tester_factory: Callable) -> None:
+        # Test with different modes (sum, mean, max)
+        indices = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
+        offsets = torch.tensor([0, 4], dtype=torch.long)
+
+        self._test_op(Model(mode='sum'), (indices, offsets), tester_factory, use_random_test_inputs=False)
+        self._test_op(Model(mode='mean'), (indices, offsets), tester_factory, use_random_test_inputs=False)
+        self._test_op(Model(mode='max'), (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_bag_padding_idx(self, tester_factory: Callable) -> None:
+        # Test with padding_idx
+        indices = torch.tensor([0, 1, 2, 0, 3, 0, 4], dtype=torch.long)
+        offsets = torch.tensor([0, 3, 6], dtype=torch.long)
+
+        self._test_op(Model(padding_idx=0), (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+        indices = torch.tensor([1, 5, 2, 5, 3, 5, 4], dtype=torch.long)
+        offsets = torch.tensor([0, 3, 6], dtype=torch.long)
+
+        self._test_op(Model(padding_idx=5), (indices, offsets), tester_factory, use_random_test_inputs=False)
+
+    def test_embedding_bag_include_last_offset(self, tester_factory: Callable) -> None:
+        # Test with include_last_offset
+        indices = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long)
+        offsets = torch.tensor([0, 4], dtype=torch.long)
+
+        self._test_op(Model(include_last_offset=True), (indices, offsets), tester_factory, use_random_test_inputs=False)
+