From 12c3e45c11766b6505235bc15e3b42f97e0eb400 Mon Sep 17 00:00:00 2001
From: Christophe Guillon <christophe.guillon@inria.fr>
Date: Fri, 12 Jul 2024 13:56:46 +0200
Subject: [PATCH] [Tensor] Fix Tensor::getIdx() method
The getIdx() method generates segfault on 0-rank coords
due to an out-of-bound access.
Also the computation was wrong for rank of coords
lower than rank of tensor by 2 or more.
---
include/aidge/data/Tensor.hpp | 12 +++++++-----
1 file changed, 7 insertions(+), 5 deletions(-)
diff --git a/include/aidge/data/Tensor.hpp b/include/aidge/data/Tensor.hpp
index 108f1f2b4..96d877ebd 100644
--- a/include/aidge/data/Tensor.hpp
+++ b/include/aidge/data/Tensor.hpp
@@ -639,6 +639,7 @@ public:
* the remaining coordinates are assumed to be 0.
* Beware: the contiguous index will only correspond to the storage index
* if the tensor is contiguous!
+ * Note that the coordIdx may be an empty vector.
*
* @param coordIdx Coordinate to an element in the tensor
* @return DimSize_t Contiguous index
@@ -646,12 +647,13 @@ public:
std::size_t getIdx(const std::vector<std::size_t>& coordIdx) const {
AIDGE_ASSERT(coordIdx.size() <= mDims.size(), "Coordinates does not match number of dimensions");
std::size_t flatIdx = 0;
- std::size_t i = 0;
- for(; i < coordIdx.size() - 1; ++i) {
- AIDGE_ASSERT(coordIdx[i] < mDims[i], "Coordinates dimensions does not fit the dimensions of the tensor");
- flatIdx = (flatIdx + coordIdx[i]) * mDims[i + 1];
+ for(std::size_t i = 0; i < mDims.size(); ++i) {
+ auto coord = i < coordIdx.size() ? coordIdx[i]: 0;
+ AIDGE_ASSERT(coord < mDims[i], "Coordinates dimensions does not fit the dimensions of the tensor");
+ auto nextDimSize = i + 1 < mDims.size() ? mDims[i + 1]: 1;
+ flatIdx = (flatIdx + coord) * nextDimSize;
}
- return flatIdx + coordIdx[i];
+ return flatIdx;
}
/**
--
GitLab