Updates to reduction and reverse operations

Simplify the reduction and reverse reference
by apply one axis at a time and default to
keep_dims. (Changes to shape can be done
separately with RESHAPE)

Signed-off-by: Dominic Symes <dominic.symes@arm.com>
Change-Id: I590bdb6bc05b1c673f86e3f45f0a43536d8f362a

2 files changed, 65 insertions, 102 deletions

diff --git a/chapters/data_layout.adoc b/chapters/data_layout.adoc
index 8f318ae..3a7c3c3 100644
--- a/chapters/data_layout.adoc
+++ b/chapters/data_layout.adoc
@@ -135,7 +135,7 @@ for (i=0; i<tensor_size(shape); i++) {
 
 ==== REVERSE
 
-Returns a tensor with the same type/values as the input, with the data reversed along the given axes. No data conversion happens during a reverse operation.
+Returns a tensor with the same type/values as the input, with the data reversed along the given axis. No data conversion happens during a reverse operation.
 
 *Arguments:*
 
@@ -143,19 +143,18 @@ 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|axes|[]|List of axes along which to reverse with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Output|out_t*|output|shape|Output tensor of same type, size as the input tensor
+|Attribute|int|axis|-|Axis to reverse
+|Output|out_t*|output|shape|Output tensor. Same shape as input tensor.
 |===
 
 *Operation Function:*
 
 [source,c]
 ----
+assert(0<=axis && axis<dimensions(shape))
 for_each (index in shape) {
     tmp_index = index;
-    for_each (axis in axes[]) {
-        tmp_index[axis] = shape[axis]-1-index[axis]
-    }
+    tmp_index[axis] = shape[axis]-1-index[axis];
     value = tensor_read<in_t>(input, shape, tmp_index);
     tensor_write<in_t>(output, shape, index, value);
 }
diff --git a/chapters/reduction.adoc b/chapters/reduction.adoc
index a9687b1..af44ab6 100644
--- a/chapters/reduction.adoc
+++ b/chapters/reduction.adoc
@@ -11,7 +11,7 @@
 
 ==== REDUCE_ALL
 
-Reduce a tensor along the given axes with a logical AND operation
+Reduce a tensor along the given axis with a logical AND operation
 
 *Arguments:*
 
@@ -19,32 +19,26 @@ Reduce a tensor along the given axes with a logical AND operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|in_shape|Input tensor from 1 to 4 dims
-|Attribute|int|axes|[]|List of axes with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Attribute|int|keep_dims|-|If 1, axes reduced will be retained with length=1.
-|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor if keep_dims set, else one less than input tensor rank.
+|Attribute|int|axis|-|Axis to reduce
+|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor.
 |===
 
 *Operation Function:*
 
 [source,c]
 ----
-tmp_shape = in_shape;
-for_each (axis in axes[]) {
-    tmp_shape[axis]=1;
-}
-if (keep_dims) assert(tmp_shape == out_shape)
-for_each (index in tmp_shape) {
-    tensor_write<in_t>(output, tmp_shape, index, true)
+assert(0<=axis && axis<dimensions(in_shape))
+assert(out_shape[axis]==1)
+for_each (index in out_shape) {
+    tensor_write<in_t>(output, out_shape, index, true)
 }
 for_each (index in in_shape) {
     tmp_index = index;
-    for each (axis in axes[]) {
-        tmp_index[axis]=0;
-    }
+    tmp_index[axis]=0;
     value = tensor_read<in_t>(input, in_shape, index)
-    acc   = tensor_read<in_t>(output, tmp_shape, tmp_index)
+    acc   = tensor_read<in_t>(output, out_shape, tmp_index)
     acc   = acc && value
-    tensor_write<in_t>(output, tmp_shape, tmp_index, acc)
+    tensor_write<in_t>(output, out_shape, tmp_index, acc)
 }
 ----
 
@@ -58,7 +52,7 @@ for_each (index in in_shape) {
 
 ==== REDUCE_ANY
 
-Reduce a tensor along the given axes with a logical OR operation
+Reduce a tensor along the given axis with a logical OR operation
 
 *Arguments:*
 
@@ -66,32 +60,26 @@ Reduce a tensor along the given axes with a logical OR operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|in_shape|Input tensor from 1 to 4 dims
-|Attribute|int|axes|[]|List of axes with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Attribute|int|keep_dims|-|If 1, axes reduced will be retained with length=1.
-|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor if keep_dims set, else one less than input tensor rank.
+|Attribute|int|axis|-|Axis to reduce
+|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor.
 |===
 
 *Operation Function:*
 
 [source,c]
 ----
-tmp_shape = in_shape;
-for_each (axis in axes[]) {
-    tmp_shape[axis]=1;
-}
-if (keep_dims) assert(tmp_shape == out_shape)
-for_each (index in tmp_shape) {
-    tensor_write<in_t>(output, tmp_shape, index, false)
+assert(0<=axis && axis<dimensions(in_shape))
+assert(out_shape[axis]==1)
+for_each (index in out_shape) {
+    tensor_write<in_t>(output, out_shape, index, false)
 }
 for_each (index in in_shape) {
     tmp_index = index;
-    for each (axis in axes[]) {
-        tmp_index[axis]=0;
-    }
+    tmp_index[axis]=0;
     value = tensor_read<in_t>(input, in_shape, index)
-    acc   = tensor_read<in_t>(output, tmp_shape, tmp_index)
+    acc   = tensor_read<in_t>(output, out_shape, tmp_index)
     acc   = acc || value
-    tensor_write<in_t>(output, tmp_shape, tmp_index, acc)
+    tensor_write<in_t>(output, out_shape, tmp_index, acc)
 }
 ----
 
@@ -105,7 +93,7 @@ for_each (index in in_shape) {
 
 ==== REDUCE_MAX
 
-Reduce a tensor along the given axes with a maximum operation
+Reduce a tensor along the given axis with a maximum operation
 
 *Arguments:*
 
@@ -113,32 +101,26 @@ Reduce a tensor along the given axes with a maximum operation
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|in_shape|Input tensor from 1 to 4 dims
-|Attribute|int|axes|[]|List of axes with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Attribute|int|keep_dims|-|If 1, axes reduced will be retained with length=1.
-|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor if keep_dims set, else one less than input tensor rank.
+|Attribute|int|axis|-|Axis to reduce
+|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor.
 |===
 
 *Operation Function:*
 
 [source,c]
 ----
-tmp_shape = in_shape;
-for_each (axis in axes[]) {
-    tmp_shape[axis]=1;
-}
-if (keep_dims) assert(tmp_shape == out_shape)
-for_each (index in tmp_shape) {
-    tensor_write<in_t>(output, tmp_shape, index, minimum<in_t>)
+assert(0<=axis && axis<dimensions(in_shape))
+assert(out_shape[axis]==1)
+for_each (index in out_shape) {
+    tensor_write<in_t>(output, out_shape, index, minimum<in_t>)
 }
 for_each (index in in_shape) {
     tmp_index = index;
-    for each (axis in axes[]) {
-        tmp_index[axis]=0;
-    }
+    tmp_index[axis]=0;
     value = tensor_read<in_t>(input, in_shape, index)
-    acc   = tensor_read<in_t>(output, tmp_shape, tmp_index)
+    acc   = tensor_read<in_t>(output, out_shape, tmp_index)
     acc   = apply_max<in_t>(acc, value)
-    tensor_write<in_t>(output, tmp_shape, tmp_index, acc)
+    tensor_write<in_t>(output, out_shape, tmp_index, acc)
 }
 ----
 
@@ -155,16 +137,15 @@ for_each (index in in_shape) {
 
 ==== REDUCE_MIN
 
-Reduce a tensor along the given axes with a minimum operation
+Reduce a tensor along the given axis with a minimum operation
 
 *Arguments:*
 |===
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|in_shape|Input tensor from 1 to 4 dims
-|Attribute|int|axes|[]|List of axes with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Attribute|int|keep_dims|-|If 1, axes reduced will be retained with length=1.
-|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor if keep_dims set, else one less than input tensor rank.
+|Attribute|int|axis|-|Axis to reduce
+|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor.
 |===
 
 *Quantization Parameters:*
@@ -175,23 +156,18 @@ Quantization is ignored when doing the REDUCE_MIN operation. The input and outpu
 
 [source,c]
 ----
-tmp_shape = in_shape;
-for_each (axis in axes[]) {
-    tmp_shape[axis]=1;
-}
-if (keep_dims) assert(tmp_shape == out_shape)
-for_each (index in tmp_shape) {
-    tensor_write<in_t>(output, tmp_shape, index, maximum<in_t>)
+assert(0<=axis && axis<dimensions(in_shape))
+assert(out_shape[axis]==1)
+for_each (index in out_shape) {
+    tensor_write<in_t>(output, out_shape, index, maximum<in_t>)
 }
 for_each (index in in_shape) {
     tmp_index = index;
-    for each (axis in axes[]) {
-        tmp_index[axis]=0;
-    }
+    tmp_index[axis]=0;
     value = tensor_read<in_t>(input, in_shape, index)
-    acc   = tensor_read<in_t>(output, tmp_shape, tmp_index)
+    acc   = tensor_read<in_t>(output, out_shape, tmp_index)
     acc   = apply_min<in_t>(acc, value)
-    tensor_write<in_t>(output, tmp_shape, tmp_index, acc)
+    tensor_write<in_t>(output, out_shape, tmp_index, acc)
 }
 ----
 
@@ -208,7 +184,7 @@ for_each (index in in_shape) {
 
 ==== REDUCE_PRODUCT
 
-Reduce a tensor along the given axes by computing the product of the axes.
+Reduce a tensor along the given axis by computing the product of the axis.
 
 *Arguments:*
 
@@ -216,32 +192,26 @@ Reduce a tensor along the given axes by computing the product of the axes.
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|in_shape|Input tensor from 1 to 4 dims
-|Attribute|int|axes|[]|List of axes with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Attribute|int|keep_dims|-|If 1, axes reduced will be retained with length=1.
-|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor if keep_dims set, else one less than input tensor rank.
+|Attribute|int|axis|-|Axis to reduce
+|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor.
 |===
 
 *Operation Function:*
 
 [source,c]
 ----
-tmp_shape = in_shape;
-for_each (axis in axes[]) {
-    tmp_shape[axis]=1;
-}
-if (keep_dims) assert(tmp_shape == out_shape)
-for_each (index in tmp_shape) {
-    tensor_write<in_t>(output, tmp_shape, index, 1.0)
+assert(0<=axis && axis<dimensions(in_shape))
+assert(out_shape[axis]==1)
+for_each (index in out_shape) {
+    tensor_write<in_t>(output, out_shape, index, 1.0)
 }
 for_each (index in in_shape) {
     tmp_index = index;
-    for each (axis in axes[]) {
-        tmp_index[axis]=0;
-    }
+    tmp_index[axis]=0;
     value = tensor_read<in_t>(input, in_shape, index)
-    acc   = tensor_read<in_t>(output, tmp_shape, tmp_index)
+    acc   = tensor_read<in_t>(output, out_shape, tmp_index)
     acc   = acc * value
-    tensor_write<in_t>(output, tmp_shape, tmp_index, acc)
+    tensor_write<in_t>(output, out_shape, tmp_index, acc)
 }
 ----
 
@@ -255,7 +225,7 @@ for_each (index in in_shape) {
 
 ==== REDUCE_SUM
 
-Reduce a tensor along the given axes by computing the sum of the axes.
+Reduce a tensor along the given axis by computing the sum of the axis.
 
 *Arguments:*
 
@@ -263,32 +233,26 @@ Reduce a tensor along the given axes by computing the sum of the axes.
 |Argument|Type|Name|Shape|Description
 
 |Input|in_t*|input|in_shape|Input tensor from 1 to 4 dims
-|Attribute|int|axes|[]|List of axes with size from 1 to rank(input1). Values are integers into the 0-based dimensions of the input tensor. An empty list is not allowed.
-|Attribute|int|keep_dims|-|If 1, axes reduced will be retained with length=1.
-|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor if keep_dims set, else one less than input tensor rank.
+|Attribute|int|axis|-|Axis to reduce
+|Output|out_t*|output|out_shape|Output tensor. Same rank as the input tensor.
 |===
 
 *Operation Function:*
 
 [source,c]
 ----
-tmp_shape = in_shape;
-for_each (axis in axes[]) {
-    tmp_shape[axis]=1;
-}
-if (keep_dims) assert(tmp_shape == out_shape)
-for_each (index in tmp_shape) {
-    tensor_write<in_t>(output, tmp_shape, index, 0)
+assert(0<=axis && axis<dimensions(in_shape))
+assert(out_shape[axis]==1)
+for_each (index in out_shape) {
+    tensor_write<in_t>(output, out_shape, index, 0)
 }
 for_each (index in in_shape) {
     tmp_index = index;
-    for each (axis in axes[]) {
-        tmp_index[axis]=0;
-    }
+    tmp_index[axis]=0;
     value = tensor_read<in_t>(input, in_shape, index)
-    acc   = tensor_read<in_t>(output, tmp_shape, tmp_index)
+    acc   = tensor_read<in_t>(output, out_shape, tmp_index)
     acc   = apply_add<in_t>(acc, value)
-    tensor_write<in_t>(output, tmp_shape, tmp_index, acc)
+    tensor_write<in_t>(output, out_shape, tmp_index, acc)
 }
 ----