commit 6986a079020ab6344c9191aa67af13beeb475593
author Rickard Bolin <rickard.bolin@arm.com> Mon Dec 19 12:33:40 2022 +0000
committer Rickard Bolin <rickard.bolin@arm.com> Wed Mar 22 10:11:21 2023 +0000
tree c5cfef45fd915dc44b313f5fd8510f375dafa806
parent 3340a88d51cdce4d00ea82a2211d9e7fd40cf5c5

MLBEDSW-6435: Implement support for ArgMax along depth dimension

- Add support for ArgMax along depth dimension with a depth limit of 127.
- Only supports 8-bit input and 32-bit output

Signed-off-by: Rickard Bolin <rickard.bolin@arm.com>
Change-Id: I5f6f0503135bebabbb1ca637f9729587b7c60740

ethosu/vela/operation.py

diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py
index 6be9dc2..6771710 100644
--- a/ethosu/vela/operation.py
+++ b/ethosu/vela/operation.py
@@ -134,7 +134,7 @@
     Add = OperatorInfo(block_type=NpuBlockType.ElementWise, indices=NNG_IFM_IFM2_INDICES)
     AddN = OperatorInfo()
     Any = OperatorInfo()
-    ArgMax = OperatorInfo()
+    ArgMax = OperatorInfo(indices=NNG_IFM_INDICES)
     ArgMin = OperatorInfo()
     AvgPool = OperatorInfo(block_type=NpuBlockType.Pooling, indices=NNG_IFM_INDICES)
     Atan2 = OperatorInfo(indices=NNG_IFM_IFM2_INDICES)