Better error checking for Argmin/Argmax (#3259)

* Better error checking for Argmin/Argmax

Argmin and Argmax are not defined if the reduction axis has length 0.

* Simplify

* Update test/type_prop/index_reduction.cpp
Co-Authored-By: Adam Procter <adam.m.procter@intel.com>

* Update test/type_prop/index_reduction.cpp
Co-Authored-By: Adam Procter <adam.m.procter@intel.com>

src/ngraph/op/util/index_reduction.cpp

@@ -91,8 +91,17 @@ void op::util::IndexReduction::validate_and_infer_types()
 
     PartialShape output_shape{PartialShape::dynamic()};
 
-    if (!rank.is_dynamic())
+    if (rank.is_static())
     {
+        Dimension d = arg_shape[m_axis];
+        if (d.is_static())
+        {
+            NODE_VALIDATION_CHECK(this,
+                                  0 != size_t(d),
+                                  "Tensor reduction axis can not be empty, shape is: ",
+                                  arg_shape);
+        }
+
         std::vector<Dimension> output_dims(size_t(rank) - 1);
         size_t j = 0;
 

test/type_prop/index_reduction.cpp

@@ -59,6 +59,44 @@ TEST(type_prop, index_reduction_invalid_rank)
     }
 }
 
+TEST(type_prop, argmin_invalid_zero_reduction_axis)
+{
+    auto a = make_shared<op::Parameter>(element::f32, Shape{2, 0});
+
+    try
+    {
+        auto argmin = make_shared<op::ArgMin>(a, 1, element::i32);
+        FAIL() << "ArgMin c-tor should throw for zero-length reduction axis";
+    }
+    catch (const NodeValidationFailure& error)
+    {
+        EXPECT_HAS_SUBSTRING(error.what(), "reduction axis can not be empty");
+    }
+    catch (...)
+    {
+        FAIL() << "Deduced type check failed for unexpected reason";
+    }
+}
+
+TEST(type_prop, argmax_invalid_zero_reduction_axis)
+{
+    auto a = make_shared<op::Parameter>(element::f32, Shape{2, 0});
+
+    try
+    {
+        auto argmax = make_shared<op::ArgMax>(a, 1, element::i32);
+        FAIL() << "ArgMax c-tor should throw for zero-length reduction axis";
+    }
+    catch (const NodeValidationFailure& error)
+    {
+        EXPECT_HAS_SUBSTRING(error.what(), "reduction axis can not be empty");
+    }
+    catch (...)
+    {
+        FAIL() << "Deduced type check failed for unexpected reason";
+    }
+}
+
 TEST(type_prop, index_reduction_invalid_index_type)
 {
     auto a = make_shared<op::Parameter>(element::f32, Shape{2, 2});