diff options
author | Michalis Spyrou <michalis.spyrou@arm.com> | 2019-01-03 11:10:25 +0000 |
---|---|---|
committer | Michalis Spyrou <michalis.spyrou@arm.com> | 2019-01-10 16:24:26 +0000 |
commit | aea14c63e2efeda9d5f7492099389d439c65204f (patch) | |
tree | 176a6181bbf00e4df078d5da0a17dd44f248958e /arm_compute/core/utils/misc/ShapeCalculator.h | |
parent | c10bc0b5db5169a6ccea02a1aaefe34f082709e5 (diff) | |
download | ComputeLibrary-aea14c63e2efeda9d5f7492099389d439c65204f.tar.gz |
COMPMID-1764 NEON: Implement ArgMax/ArgMin
Change-Id: Ibe23aa90b36ffd8553d1d1c35fada5d300fab829
Reviewed-on: https://review.mlplatform.org/475
Reviewed-by: Isabella Gottardi <isabella.gottardi@arm.com>
Tested-by: Arm Jenkins <bsgcomp@arm.com>
Reviewed-by: Giuseppe Rossini <giuseppe.rossini@arm.com>
Diffstat (limited to 'arm_compute/core/utils/misc/ShapeCalculator.h')
-rw-r--r-- | arm_compute/core/utils/misc/ShapeCalculator.h | 15 |
1 files changed, 15 insertions, 0 deletions
diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h index 70727424b3..619234d306 100644 --- a/arm_compute/core/utils/misc/ShapeCalculator.h +++ b/arm_compute/core/utils/misc/ShapeCalculator.h @@ -1033,6 +1033,21 @@ inline TensorShape compute_tiled_shape(const TensorShape &input_shape, const Mul return tiled_shape; } +/** Calculate the reduced shape of a tensor given an axis + * + * @param[in] input Input tensor info + * @param[in] axis Axis on which to perform reduction + * + * @return the calculated shape + */ +inline TensorShape compute_reduced_shape(const TensorShape &input, unsigned int axis) +{ + TensorShape output_shape{ input }; + output_shape.set(axis, 1); + + return output_shape; +} + /** Calculate the upsampled shape of a tensor * * @param[in] input Input tensor info |