diff options
Diffstat (limited to 'chapters/pseudocode.adoc')
-rw-r--r-- | chapters/pseudocode.adoc | 6 |
1 files changed, 6 insertions, 0 deletions
diff --git a/chapters/pseudocode.adoc b/chapters/pseudocode.adoc index 146b5d7..c026089 100644 --- a/chapters/pseudocode.adoc +++ b/chapters/pseudocode.adoc @@ -94,6 +94,12 @@ size_t tensor_size(dim_t shape) { } 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(dim_t shape, int axis) { + return (axis >= rank(shape)) ? 1 : shape[axis]; +} ---- ==== Tensor Read |