Note limits on axis parameter

This also adds a restriction on CONCAT to have all inputs be the
same rank.

Update the ARGMAX pseudocode to handle corner cases properly and
use axis/rank consistent with other operators.

Signed-off-by: Eric Kunze <eric.kunze@arm.com>
Change-Id: I941ac172ee79424ac04d562cedebb75da76659a5

3 files changed, 34 insertions, 24 deletions

diff --git a/chapters/data_layout.adoc b/chapters/data_layout.adoc
index 4b01d55..de90322 100644
--- a/chapters/data_layout.adoc
+++ b/chapters/data_layout.adoc
@@ -18,8 +18,8 @@ No data conversion happens during a concat operation.
 |===
 |Argument|Type|Name|Shape|Description
 
-|Input|in_t**|input1|shapes1[]|List of input tensors
-|Attribute|int|axis|-|Axis along which concatenation is to occur
+|Input|in_t**|input1|shapes1[]|List of input tensors. All inputs must have the same rank
+|Attribute|int|axis|-|Axis along which concatenation is to occur, in range from 0 to rank(shape)-1
 |Output|in_t*|output|shape|Output tensor
 |===
 
@@ -149,7 +149,7 @@ Returns a tensor with the same type/values as the input, with the data reversed
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape|Input tensor from 1 to 4 dims
-|Attribute|int|axis|-|Axis to reverse
+|Attribute|int|axis|-|Axis to reverse, in range from 0 to rank(shape)-1
 |Output|in_t*|output|shape|Output tensor. Same shape as input tensor.
 |===
 
diff --git a/chapters/reduction.adoc b/chapters/reduction.adoc
index e605386..b84b7f0 100644
--- a/chapters/reduction.adoc
+++ b/chapters/reduction.adoc
@@ -19,7 +19,7 @@ Reduce a tensor along the given axis with a logical AND operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
-|Attribute|int32_t|axis|-|Axis to reduce
+|Attribute|int32_t|axis|-|Axis to reduce, in range from 0 to rank(shape1)-1
 |Output|in_t*|output|shape|Output tensor. Same rank as the input tensor.
 |===
 
@@ -60,7 +60,7 @@ Reduce a tensor along the given axis with a logical OR operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
-|Attribute|int32_t|axis|-|Axis to reduce
+|Attribute|int32_t|axis|-|Axis to reduce, in range from 0 to rank(shape1)-1
 |Output|in_t*|output|shape|Output tensor. Same rank as the input tensor.
 |===
 
@@ -101,7 +101,7 @@ Reduce a tensor along the given axis with a maximum operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
-|Attribute|int32_t|axis|-|Axis to reduce
+|Attribute|int32_t|axis|-|Axis to reduce, in range from 0 to rank(shape1)-1
 |Output|in_t*|output|shape|Output tensor. Same rank as the input tensor.
 |===
 
@@ -144,7 +144,7 @@ Reduce a tensor along the given axis with a minimum operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
-|Attribute|int32_t|axis|-|Axis to reduce
+|Attribute|int32_t|axis|-|Axis to reduce, in range from 0 to rank(shape1)-1
 |Output|in_t*|output|shape|Output tensor. Same rank as the input tensor.
 |===
 
@@ -192,7 +192,7 @@ Reduce a tensor along the given axis by computing the product of the axis.
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
-|Attribute|int32_t|axis|-|Axis to reduce
+|Attribute|int32_t|axis|-|Axis to reduce, in range from 0 to rank(shape1)-1
 |Output|in_t*|output|shape|Output tensor. Same rank as the input tensor.
 |===
 
@@ -233,7 +233,7 @@ Reduce a tensor along the given axis by computing the sum of the axis.
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
-|Attribute|int32_t|axis|-|Axis to reduce
+|Attribute|int32_t|axis|-|Axis to reduce, in range from 0 to rank(shape1)-1
 |Output|in_t*|output|shape|Output tensor. Same rank as the input tensor.
 |===
 
diff --git a/chapters/tensor_ops.adoc b/chapters/tensor_ops.adoc
index 291751f..b2c220e 100644
--- a/chapters/tensor_ops.adoc
+++ b/chapters/tensor_ops.adoc
@@ -18,9 +18,9 @@ This returns the index with the largest value across the given axis of the input
 |===
 |Argument|Type|Name|Shape|Description
 
-|Input|in_t*|input|shape1|Input tensor dimension k \<=4
-|Attribute|int|axis|-|Axis in range 0 to k-1
-|Output|out_t*|output|shape|Output tensor dimension k-1
+|Input|in_t*|input|shape1|Input tensor with rank from 1 to 4
+|Attribute|int|axis|-|Axis in range from 0 to rank(shape1)-1
+|Output|out_t*|output|shape|Output tensor, with rank = rank(shape1)-1
 |===
 
 *Quantization Parameters:*
@@ -31,20 +31,30 @@ None
 
 [source,c++]
 ----
-assert(axis >= 0 && axis < k && k <=4);
-left_shape = shape1[0:axis-1];
-right_shape = shape1[axis+1:k-1];
+assert(axis >= 0 && axis < rank(shape1) && rank(shape1) <= 4);
+if (axis == 0) {
+    left_shape = [];
+} else {
+    left_shape = shape1[0:axis - 1];
+}
+if (axis == rank(shape1)-1) {
+    right_shape = [];
+} else {
+    right_shape = shape1[axis+1:rank(shape1) - 1];
+}
 assert(flatten(left_shape, right_shape) == shape);
-for_each(left_index in left_shape, right_index in right_shape )
-    in_t max_value = minimum_value<in_t>;
-    int32_t max_index = 0;
-    for (i = 0; i < shape[axis]; i++) {
-        index = flatten(left_index, [i], right_index);
-        in_t value = tensor_read<in_t>(input, shape1, index);
-        if (value > max_value) { max_value = value; max_index=i; }
+for_each(left_index in left_shape) {
+    for_each(right_index in right_shape) {
+        in_t max_value = minimum_value<in_t>;
+        int32_t max_index = 0;
+        for (i = 0; i < shape[axis]; i++) {
+            index = flatten(left_index, [i], right_index);
+            in_t value = tensor_read<in_t>(input, shape1, index);
+            if (value > max_value) { max_value = value; max_index = i; }
+        }
+        index = flatten(left_index, right_index);
+        tensor_write<int32_t>(output, shape, index, max_index);
     }
-    index = flatten(left_index, right_index);
-    tensor_write<int32_t>(output, shape, index, max_index);
 }
 ----