diff options
Diffstat (limited to 'arm_compute/core/utils/misc/ShapeCalculator.h')
| -rw-r--r-- | arm_compute/core/utils/misc/ShapeCalculator.h | 18 |
1 files changed, 18 insertions, 0 deletions
diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h index 09f558d8b0..e0b6ce639f 100644 --- a/arm_compute/core/utils/misc/ShapeCalculator.h +++ b/arm_compute/core/utils/misc/ShapeCalculator.h @@ -28,6 +28,8 @@ #include "arm_compute/core/ITensorInfo.h" #include "arm_compute/core/Utils.h" +#include "arm_compute/core/utils/helpers/tensor_transform.h" + #include <cmath> namespace arm_compute @@ -432,6 +434,22 @@ inline TensorShape compute_mm_shape(const ITensorInfo &input0, const ITensorInfo return output_shape; } +inline TensorShape compute_strided_slice_shape(const ITensorInfo &input, + const Coordinates &starts, const Coordinates &ends, const Coordinates &strides, + int32_t begin_mask, int32_t end_mask, int32_t shrink_axis_mask) +{ + using namespace arm_compute::helpers::tensor_transform; + + const TensorShape &input_shape = input.tensor_shape(); + + // Get actual start, end coordinates and strides + const Coordinates final_strides = strided_slice_strides(input_shape, strides); + const Coordinates starts_abs = strided_slice_absolute_start_coords(input_shape, starts, final_strides, begin_mask); + const Coordinates ends_abs = strided_slice_absolute_end_coords(input_shape, starts_abs, ends, final_strides, end_mask, shrink_axis_mask); + + return compute_strided_slice_output_shape(input_shape, starts_abs, ends_abs, final_strides); +} + template <typename T> inline TensorShape extract_shape(T *data) { |