Add new ops ARGMAX and ARGMIN

Note that I could not name it ARG_MAX because of a name collision with a
macro in libc/kernel/uapi/linux/limits.h.

Bug: 113558658
Test: NeuralNetworksTest_static
Change-Id: Ibb1cd619d1a5428ab4bab63412b21f5a63ac2cd5
Merged-In: Ibb1cd619d1a5428ab4bab63412b21f5a63ac2cd5
(cherry picked from commit cc700116c959e4d2f4a0f1f2bfd9fa13824339d7)

common/OperationsUtils.cpp

diff --git a/common/OperationsUtils.cpp b/common/OperationsUtils.cpp
index 5371718..e918ae3 100644
--- a/common/OperationsUtils.cpp
+++ b/common/OperationsUtils.cpp
@@ -53,6 +53,17 @@
     return count;
 }
 
+uint32_t getNumberOfElements(const Shape& shape,
+                             size_t firstAxisInclusive,
+                             size_t lastAxisExclusive) {
+    NN_CHECK(lastAxisExclusive <= shape.dimensions.size());
+    uint32_t count = 1;
+    for (size_t i = firstAxisInclusive; i < lastAxisExclusive; i++) {
+        count *= shape.dimensions[i];
+    }
+    return count;
+}
+
 uint32_t getNumberOfDimensions(const Shape& shape) {
     return shape.dimensions.size();
 }
@@ -880,5 +891,23 @@
 
     return true;
 }
+
+bool argMinMaxPrepare(const Shape& input, int32_t axis, Shape* output) {
+    NN_OPS_CHECK(0 <= axis && axis < getNumberOfDimensions(input));
+
+    output->type = OperandType::TENSOR_INT32;
+
+    // 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());
+
+    return true;
+}
 } // namespace nn
 } // namespace android