diff options
| author | Rickard Bolin <rickard.bolin@arm.com> | 2022-12-19 12:33:40 +0000 |
|---|---|---|
| committer | Rickard Bolin <rickard.bolin@arm.com> | 2023-03-22 10:11:21 +0000 |
| commit | 6986a079020ab6344c9191aa67af13beeb475593 (patch) | |
| tree | c5cfef45fd915dc44b313f5fd8510f375dafa806 /ethosu/vela/operation.py | |
| parent | 3340a88d51cdce4d00ea82a2211d9e7fd40cf5c5 (diff) | |
| download | ethos-u-vela-6986a079020ab6344c9191aa67af13beeb475593.tar.gz | |
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
Diffstat (limited to 'ethosu/vela/operation.py')
| -rw-r--r-- | ethosu/vela/operation.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/ethosu/vela/operation.py b/ethosu/vela/operation.py index 6be9dc25..67717104 100644 --- a/ethosu/vela/operation.py +++ b/ethosu/vela/operation.py @@ -134,7 +134,7 @@ class Op(Enum): 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) |