diff options
Diffstat (limited to 'pseudocode/library/tensor_utils.tosac')
-rw-r--r-- | pseudocode/library/tensor_utils.tosac | 47 |
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]; +} |