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

2 files changed, 15 insertions, 12 deletions

diff --git a/ethosu/vela/reader_util.py b/ethosu/vela/reader_util.py
index b5a058e3..8ca1c156 100644
--- a/ethosu/vela/reader_util.py
+++ b/ethosu/vela/reader_util.py
@@ -15,8 +15,6 @@
 # limitations under the License.
 # Description:
 # Utlity function for reading .tosa and .tflite files
-import numpy as np
-
 from .operation import Op
 from .operation import Operation
 
@@ -30,20 +28,14 @@ def decode_str(s):
 def clone_and_reshape_tensor(src_tens, reorder, set_unique):
     tens = src_tens.clone("_reshape", set_unique)
     if reorder is None:
-        # 1D shape requested
-        tens.shape = [np.prod(src_tens.shape)]
+        # reorder of None is a special case meaning 1D shape requested
+        tens.as_1D()
     else:
-        tens.shape = [src_tens.shape[idx] for idx in reorder]
+        tens.transpose(reorder)
+
     tens.bandwidth_shape = tens.shape
     tens.storage_shape = tens.shape
 
-    if tens.values is not None:
-        if reorder is None:
-            # 1D shape requested
-            tens.values = tens.values.reshape(tens.shape)
-        else:
-            tens.values = tens.values.transpose(reorder)
-
     op = Operation(Op.Const, tens.name)
     op.set_output_tensor(tens)
     return tens
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