diff options
Diffstat (limited to 'libs/tflite/src/TensorShapeUtils.cpp')
-rw-r--r-- | libs/tflite/src/TensorShapeUtils.cpp | 48 |
1 files changed, 48 insertions, 0 deletions
diff --git a/libs/tflite/src/TensorShapeUtils.cpp b/libs/tflite/src/TensorShapeUtils.cpp new file mode 100644 index 000000000..b5d906719 --- /dev/null +++ b/libs/tflite/src/TensorShapeUtils.cpp @@ -0,0 +1,48 @@ +#include "tflite/TensorShapeUtils.h" + +namespace nnfw +{ +namespace tflite +{ + +nnfw::misc::tensor::Shape broadcast(const nnfw::misc::tensor::Shape &lhs_shape, + const nnfw::misc::tensor::Shape &rhs_shape) +{ + const uint32_t lhs_rank = lhs_shape.rank(); + const uint32_t rhs_rank = rhs_shape.rank(); + const uint32_t out_rank = std::max(lhs_rank, rhs_rank); + + // TODO Simplify implementation + std::vector<int32_t> lhs_normalized_dims; + std::vector<int32_t> rhs_normalized_dims; + + for (uint32_t n = 0; n < out_rank - lhs_rank; ++n) + { + lhs_normalized_dims.emplace_back(1); + } + for (uint32_t axis = 0; axis < lhs_rank; ++axis) + { + lhs_normalized_dims.emplace_back(lhs_shape.dim(axis)); + } + + for (uint32_t n = 0; n < out_rank - rhs_rank; ++n) + { + rhs_normalized_dims.emplace_back(1); + } + for (uint32_t axis = 0; axis < rhs_rank; ++axis) + { + rhs_normalized_dims.emplace_back(rhs_shape.dim(axis)); + } + + nnfw::misc::tensor::Shape out_shape(out_rank); + + for (uint32_t axis = 0; axis < out_rank; ++axis) + { + out_shape.dim(axis) = std::max(lhs_normalized_dims.at(axis), rhs_normalized_dims.at(axis)); + } + + return out_shape; +} + +} // namespace tflite +} // namespace nnfw |