1 files changed, 47 insertions, 0 deletions
diff --git a/pseudocode/library/tensor_utils.tosac b/pseudocode/library/tensor_utils.tosac
new file mode 100644
index 0000000..e8c7c6d
--- /dev/null
+++ b/pseudocode/library/tensor_utils.tosac
@@ -0,0 +1,47 @@
+//
+// This confidential and proprietary software may be used only as
+// authorised by a licensing agreement from ARM Limited
+// (C) COPYRIGHT 2020-2024 ARM Limited
+// ALL RIGHTS RESERVED
+// The entire notice above must be reproduced on all authorised
+// copies and copies may only be made to the extent permitted
+// by a licensing agreement from ARM Limited.
+
+// Convert tensor index coordinates to an element offset
+size_t tensor_index_to_offset(shape_t shape, shape_t index) {
+    size_t size = tensor_size(shape);  // check tensor shape is valid
+    size_t offset = 0;
+    for (int32_t i = 0; i < rank(shape); i++) {
+        REQUIRE(index[i] >= 0 && index[i] < shape[i]);
+        offset = offset * shape[i] + index[i];
+    }
+    return offset;
+}
+
+// Convert an element offset to tensor index coordinates
+shape_t tensor_offset_to_index(shape_t shape, size_t offset) {
+    size_t size = tensor_size(shape);  // check tensor shape is valid
+    REQUIRE(offset < size);
+    shape_t index(rank(shape));    // index has rank(shape) indicies
+    for(int32_t i = rank(shape) - 1; i >= 0; i--) {
+        index[i] = offset % shape[i];
+        offset /= shape[i];
+    }
+    return index;
+}
+
+// Check the tensor shape is valid and return the tensor size in elements
+size_t tensor_size(shape_t shape) {
+    size_t size = 1;
+    for (int32_t i = 0; i < rank(shape); i++) {
+        REQUIRE(1 <= shape[i] && shape[i] <= maximum<size_t> / size);
+        size *= shape[i];
+    }
+    return size;
+}
+
+// Return the size of the tensor in the given axis
+// For a rank=0 tensor, returns 1 for all axes
+size_t shape_dim(shape_t shape, int axis) {
+    return (axis >= rank(shape)) ? 1 : shape[axis];
+}