diff options
Diffstat (limited to 'ethosu/vela/reader_util.py')
| -rw-r--r-- | ethosu/vela/reader_util.py | 14 |
1 files changed, 12 insertions, 2 deletions
diff --git a/ethosu/vela/reader_util.py b/ethosu/vela/reader_util.py index 476b70aa..b5a058e3 100644 --- a/ethosu/vela/reader_util.py +++ b/ethosu/vela/reader_util.py @@ -15,6 +15,8 @@ # 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 @@ -27,12 +29,20 @@ def decode_str(s): def clone_and_reshape_tensor(src_tens, reorder, set_unique): tens = src_tens.clone("_reshape", set_unique) - tens.shape = [src_tens.shape[idx] for idx in reorder] + if reorder is None: + # 1D shape requested + tens.shape = [np.prod(src_tens.shape)] + else: + tens.shape = [src_tens.shape[idx] for idx in reorder] tens.bandwidth_shape = tens.shape tens.storage_shape = tens.shape if tens.values is not None: - tens.values = tens.values.transpose(reorder) + 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) |