IVGCVSW-2749 Throw exception in TensorShape when requested index >= number of dimensions

Change-Id: I3589b1e901b0f81f6bb17848046a22829f91bb9e
Signed-off-by: Aron Virginas-Tar <Aron.Virginas-Tar@arm.com>
diff --git a/include/armnn/Tensor.hpp b/include/armnn/Tensor.hpp
index f4d7f9f..9380a96 100644
--- a/include/armnn/Tensor.hpp
+++ b/include/armnn/Tensor.hpp
@@ -31,15 +31,9 @@
 
     TensorShape& operator=(const TensorShape& other);
 
-    unsigned int operator[](unsigned int i) const
-    {
-        return m_Dimensions.at(i);
-    }
+    unsigned int operator[](unsigned int i) const;
 
-    unsigned int& operator[](unsigned int i)
-    {
-        return m_Dimensions.at(i);
-    }
+    unsigned int& operator[](unsigned int i);
 
     bool operator==(const TensorShape& other) const;
     bool operator!=(const TensorShape& other) const;
@@ -50,6 +44,8 @@
 private:
     std::array<unsigned int, MaxNumOfTensorDimensions> m_Dimensions;
     unsigned int m_NumDimensions;
+
+    void CheckDimensionIndex(unsigned int i) const;
 };
 
 class TensorInfo
diff --git a/src/armnn/Tensor.cpp b/src/armnn/Tensor.cpp
index 6e09e3b..da19e5b 100644
--- a/src/armnn/Tensor.cpp
+++ b/src/armnn/Tensor.cpp
@@ -11,6 +11,8 @@
 #include <boost/log/trivial.hpp>
 #include <boost/numeric/conversion/cast.hpp>
 
+#include <sstream>
+
 namespace armnn
 {
 
@@ -78,6 +80,18 @@
     return *this;
 }
 
+unsigned int TensorShape::operator[](unsigned int i) const
+{
+    CheckDimensionIndex(i);
+    return m_Dimensions.at(i);
+}
+
+unsigned int& TensorShape::operator[](unsigned int i)
+{
+    CheckDimensionIndex(i);
+    return m_Dimensions.at(i);
+}
+
 bool TensorShape::operator==(const TensorShape& other) const
 {
     return ((m_NumDimensions == other.m_NumDimensions) &&
@@ -105,6 +119,16 @@
     return count;
 }
 
+void TensorShape::CheckDimensionIndex(unsigned int i) const
+{
+    if (i >= m_NumDimensions)
+    {
+        std::stringstream errorMessage;
+        errorMessage << "Invalid dimension index: " << i << " (number of dimensions is " << m_NumDimensions << ")";
+        throw InvalidArgumentException(errorMessage.str(), CHECK_LOCATION());
+    }
+}
+
 // ---
 // --- TensorInfo
 // ---