MLBEDSW-7082: Weight tensor shape is incorrect

 - A bug was introduced by using the original_shape attribute that
causes CPU CONV2D ops to fail to run due to an incorrect weight
tensor shape
 - This was due to the original_shape not being modified when a
transpose was performed on the weight tensor
 - The fix was to transpose the original_shape just like the current
shape

Signed-off-by: Tim Hall <tim.hall@arm.com>
Change-Id: Ied72316463d26c502cf931b9dd5784041c42ab66

1 files changed, 11 insertions, 0 deletions

diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py
index 673208ac..8e553e82 100644
--- a/ethosu/vela/tensor.py
+++ b/ethosu/vela/tensor.py
@@ -484,6 +484,17 @@ class Tensor:
         res.src_tensor = self
         return res
 
+    def as_1D(self):
+        self.shape = [np.prod(self.shape)]
+        if self.values is not None:
+            self.values = self.values.reshape(self.shape)
+
+    def transpose(self, reorder):
+        self.shape = [self.shape[idx] for idx in reorder]
+        self._original_shape = [self._original_shape[idx] for idx in reorder]
+        if self.values is not None:
+            self.values = self.values.transpose(reorder)
+
     def copy_compressed_weight_info(self, src_tens: "Tensor"):
         # Copies compressed values + all related weight compression info from the given tensor
         self.equivalence_id = src_tens.equivalence_id