aboutsummaryrefslogtreecommitdiff
path: root/chapters/pseudocode.adoc
diff options
context:
space:
mode:
Diffstat (limited to 'chapters/pseudocode.adoc')
-rw-r--r--chapters/pseudocode.adoc6
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