aboutsummaryrefslogtreecommitdiff
path: root/arm_compute/core/utils/misc/ShapeCalculator.h
diff options
context:
space:
mode:
Diffstat (limited to 'arm_compute/core/utils/misc/ShapeCalculator.h')
-rw-r--r--arm_compute/core/utils/misc/ShapeCalculator.h22
1 files changed, 11 insertions, 11 deletions
diff --git a/arm_compute/core/utils/misc/ShapeCalculator.h b/arm_compute/core/utils/misc/ShapeCalculator.h
index f2f5a30b6a..72b7675749 100644
--- a/arm_compute/core/utils/misc/ShapeCalculator.h
+++ b/arm_compute/core/utils/misc/ShapeCalculator.h
@@ -1083,24 +1083,24 @@ inline TensorShape compute_batch_to_space_shape(const ITensorInfo *input, const
/** Calculate the depth to space output shape of a tensor
*
- * @param[in] input Input tensor info
- * @param[in] block Block shape value
+ * @param[in] input_shape Input tensor shape
+ * @param[in] data_layout Operation data layout
+ * @param[in] block Block shape value
*
* @return the calculated shape
*/
-inline TensorShape compute_depth_to_space_shape(const ITensorInfo *input, int block)
+inline TensorShape compute_depth_to_space_shape(const TensorShape &input_shape, DataLayout data_layout, int block)
{
ARM_COMPUTE_ERROR_ON(block < 2);
- const DataLayout data_layout = input->data_layout();
- const int idx_width = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
- const int idx_height = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
- const int idx_channel = get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL);
+ const int idx_width = get_data_layout_dimension_index(data_layout, DataLayoutDimension::WIDTH);
+ const int idx_height = get_data_layout_dimension_index(data_layout, DataLayoutDimension::HEIGHT);
+ const int idx_channel = get_data_layout_dimension_index(data_layout, DataLayoutDimension::CHANNEL);
- TensorShape output_shape{ input->tensor_shape() };
- output_shape.set(idx_width, input->dimension(idx_width) * block);
- output_shape.set(idx_height, input->dimension(idx_height) * block);
- output_shape.set(idx_channel, input->dimension(idx_channel) / (block * block));
+ TensorShape output_shape{ input_shape };
+ output_shape.set(idx_width, input_shape[idx_width] * block);
+ output_shape.set(idx_height, input_shape[idx_height] * block);
+ output_shape.set(idx_channel, input_shape[idx_channel] / (block * block));
return output_shape;
}