Fix null ptr dereference in cpu implementation of operations

Operations fixed:
* MEAN
* ARGMIN/ARGMAX
* STRIDED_SLICE

The operations would crash when provided with inputs that resulted in an
empty output shape.

The change fixes the bug by making the operations output a tensor of
size [1] in this case.

Also, update the documentation to clarify this behaviour and move
squeeze operation test to appropriate spec directory.

Bug: 155508675
Bug: 155660285
Bug: 155508675
Bug: 155238914
Test: NNTest_static
Change-Id: Ia865c26021dd4d781659957049dd567beeaeae99
Merged-In: Ia865c26021dd4d781659957049dd567beeaeae99
(cherry picked from commit 4a9c8972e3a671e9908d14fd46486ea5dc584a84)

2 files changed, 18 insertions, 5 deletions

diff --git a/nn/common/OperationsUtils.cpp b/nn/common/OperationsUtils.cpp
index e5031473e..e8dd3e2ba 100644
--- a/nn/common/OperationsUtils.cpp
+++ b/nn/common/OperationsUtils.cpp
@@ -658,6 +658,10 @@ bool meanPrepare(const Shape& input, const int32_t* axisData, const Shape& axisS
                 outDims[idx - numSkipAxis] = getSizeOfDimension(input, idx);
             }
         }
+        // Handle the case when all dimensions are removed
+        if (outDims.empty()) {
+            outDims.push_back(1);
+        }
         output->dimensions = outDims;
     }
 
@@ -675,11 +679,15 @@ bool argMinMaxPrepare(const Shape& input, int32_t axis, Shape* output) {
 
     // Copy the input dimensions, omitting the axis dimension.
     output->dimensions.clear();
-    output->dimensions.reserve(getNumberOfDimensions(input) - 1);
-    output->dimensions.insert(output->dimensions.end(), input.dimensions.begin(),
-                              input.dimensions.begin() + axis);
-    output->dimensions.insert(output->dimensions.end(), input.dimensions.begin() + axis + 1,
-                              input.dimensions.end());
+    if (getNumberOfDimensions(input) > 1) {
+        output->dimensions.reserve(getNumberOfDimensions(input) - 1);
+        output->dimensions.insert(output->dimensions.end(), input.dimensions.begin(),
+                                  input.dimensions.begin() + axis);
+        output->dimensions.insert(output->dimensions.end(), input.dimensions.begin() + axis + 1,
+                                  input.dimensions.end());
+    } else {
+        output->dimensions.push_back(1);
+    }
 
     return true;
 }
diff --git a/nn/common/operations/StridedSlice.cpp b/nn/common/operations/StridedSlice.cpp
index cd972d3de..5ff5aeca8 100644
--- a/nn/common/operations/StridedSlice.cpp
+++ b/nn/common/operations/StridedSlice.cpp
@@ -191,6 +191,11 @@ bool prepare(IOperationExecutionContext* context) {
         }
     }
 
+    // Handle the case when all dimensions are removed
+    if (outDims.empty()) {
+        outDims.push_back(1);
+    }
+
     Shape outputShape = context->getOutputShape(kOutputTensor);
     NN_RET_CHECK(SetShape(inputShape, &outputShape));
     outputShape.dimensions = outDims;