diff options
Diffstat (limited to 'arm_compute/core/ITensorInfo.h')
-rw-r--r-- | arm_compute/core/ITensorInfo.h | 46 |
1 files changed, 46 insertions, 0 deletions
diff --git a/arm_compute/core/ITensorInfo.h b/arm_compute/core/ITensorInfo.h index 9112f3ea18..b5677dffd6 100644 --- a/arm_compute/core/ITensorInfo.h +++ b/arm_compute/core/ITensorInfo.h @@ -30,6 +30,7 @@ #include "arm_compute/core/Types.h" #include "arm_compute/core/Utils.h" #include "arm_compute/core/utils/misc/ICloneable.h" +#include "arm_compute/core/utils/misc/utility.h" #include <cstddef> @@ -221,6 +222,51 @@ public: * @return A QuantizationInfo containing the scale and offset. */ virtual QuantizationInfo quantization_info() const = 0; + + /** If infos are broadcast compatible tensor info's, return the broadcasted shape and the intersection of + * the broadcasted valid regions of the tensors. + * + * Two tensor info's are broadcast compatible if their shapes are broadcast compatible. + * + * Two tensor shapes are broadcast compatible if for each dimension, they're equal or one of them is 1. + * + * If two shapes are compatible, each dimension in the broadcasted shape is the max of the original dimensions. + * + * @param[in] infos Tensor info's. + * + * @return The broadcasted shape and valid region, or an empty shape and valid region if the info's are + * not broadcast compatible. + */ + template <typename... Infos> + static std::pair<TensorShape, ValidRegion> broadcast_shape_and_valid_region(const Infos &... infos) + { + TensorShape bc_shape = TensorShape::broadcast_shape(infos.tensor_shape()...); + ValidRegion bc_valid_region{ Coordinates(), bc_shape }; + + auto broadcast_valid_region = [&bc_valid_region](const ITensorInfo & info) + { + if(info.num_dimensions() != 0) + { + for(size_t d = 0; d < bc_valid_region.shape.num_dimensions(); ++d) + { + const bool is_broadcast = (info.tensor_shape()[d] == 1); + + const int anchor_max = std::max(bc_valid_region.anchor[d], info.valid_region().anchor[d]); + const size_t valid_min = std::min(bc_valid_region.shape[d], info.valid_region().shape[d]); + + if(!is_broadcast || (valid_min == 0)) + { + bc_valid_region.anchor.set(d, anchor_max); + bc_valid_region.shape.set(d, valid_min); + } + } + } + }; + + utility::for_each(broadcast_valid_region, infos...); + + return std::pair<TensorShape, ValidRegion>(bc_shape, bc_valid_region); + } }; } #endif /*__ARM_COMPUTE_TENSORINFO_H__ */ |