From 8e9267d8873e67ad3d5068f5c85c6f8f8cf0b79d Mon Sep 17 00:00:00 2001 From: Michiel Olieslagers Date: Tue, 26 Nov 2024 14:32:16 +0000 Subject: [PATCH] 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 --- .../quantizer/quantization_annotation/generic_annotator.py | 1 + backends/arm/test/quantizer/test_generic_annotater.py | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/backends/arm/quantizer/quantization_annotation/generic_annotator.py b/backends/arm/quantizer/quantization_annotation/generic_annotator.py index fe8c9650e7..5db29c95e8 100644 --- 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 index 3d39463a42..353c8b6019 100644 --- 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 + ), + )