Added support for torch.concat operator.

Added torch.ops.aten.concatenate.default in the list of supported operators in the generic annotator. Also added a test to check the functionality of the newly added operator. Change-Id: I24d0fbf96bf9fd57ebd059cefd2af0f525fe8235
pytorch · Dec 3, 2024 · 8e9267d · 8e9267d
1 parent 4f47cc9
commit 8e9267d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/backends/arm/quantizer/quantization_annotation/generic_annotator.py b/backends/arm/quantizer/quantization_annotation/generic_annotator.py
@@ -53,6 +53,7 @@
     torch.ops.aten.tile.default,
     torch.ops.aten.flip.default,
     torch.ops.aten.cat.default,
+    torch.ops.aten.concatenate.default,
     torch.ops.aten.stack.default,
     torch.ops.aten.chunk.default,
     torch.ops.aten.contiguous.default,

diff --git a/backends/arm/test/quantizer/test_generic_annotater.py b/backends/arm/test/quantizer/test_generic_annotater.py
@@ -86,3 +86,10 @@ def test_flip(self):
         self.check_annotation(
             SingleOpModel(torch.flip, (torch.randn(2, 4),), dims=(0, 1)),
         )
+
+    def test_concat(self):
+        self.check_annotation(
+            SingleOpModel(
+                torch.concatenate, ((torch.randn(2, 3), torch.randn(2, 3)),), dim=0
+            ),
+        )