Fix CUDA reduction ops handling of optional axes input

onnxruntime/core/providers/coreml/builders/impl/reduction_op_builder.cc

@@ -89,20 +89,29 @@
                                            const logging::Logger& logger) const {
   const auto& input_defs = node.InputDefs();
 
+  NodeAttrHelper helper(node);
+
+  // noop_with_empty_axes defaults to false and is only available in newer opsets where 'axes' is an optional input
+  // so we don't need to check the 'axes' attribute used in older opsets here.
+  const bool noop_with_empty_axes = helper.Get("noop_with_empty_axes", 0) != 0;
+  bool empty_axes = true;
+
   if (input_defs.size() > 1 && input_defs[1]->Exists()) {
+    // 'axes' is optional input in new opsets
     const auto& axes_name = input_defs[1]->Name();
     const auto& initializers = input_params.graph_viewer.GetAllInitializedTensors();
     if (!Contains(initializers, axes_name)) {
       LOGS(logger, VERBOSE) << "Axes of reduction must be a constant initializer";
       return false;
     }
 
-    NodeAttrHelper helper(node);
+    empty_axes = initializers.at(axes_name)->int64_data_size() == 0;
+  }
 
-    if (initializers.at(axes_name)->int64_data_size() == 0 && helper.Get("noop_with_empty_axes", 0) != 0) {
-      LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers" << std::endl;
-      return false;
-    }
+  if (empty_axes && noop_with_empty_axes) {
+    // TODO: When we add ML Program support we should enable this as it makes the node an Identity op
+    LOGS(logger, VERBOSE) << "CoreML doesn't support noop on empty axes for reduction layers" << std::endl;
+    return false;
   }
 
   return true;

onnxruntime/core/providers/cuda/reduction/reduction_ops.cc

@@ -645,10 +645,9 @@ Status ReduceKernel<allow_multi_axes>::ComputeImpl(OpKernelContext* ctx, cudnnRe
   TensorShapeVector axes;
 
   size_t num_inputs = ctx->InputCount();
-  if (num_inputs == 2) {
+  const Tensor* axes_tensor = num_inputs == 2 ? ctx->Input<Tensor>(1) : nullptr;  // optional input. may be nullptr.
+  if (axes_tensor != nullptr) {
     // override the attribute value with the input value for reduction_axes
-    const Tensor* axes_tensor = ctx->Input<Tensor>(1);
-    ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null");
     ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor.");
     auto nDims = static_cast<size_t>(axes_tensor->Shape()[0]);
     const auto* data = axes_tensor->Data<int64_t>();

onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc

@@ -2780,6 +2780,22 @@ TEST(ReductionOpTest, ReduceSum_empty_axes_input_initializer_opset_13) {
   test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider});
 }
 
+TEST(ReductionOpTest, ReduceSum_missing_axes_input_noop_opset_13) {
+  OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain);
+  test.AddAttribute("keepdims", (int64_t)0);
+  test.AddAttribute("noop_with_empty_axes", (int64_t)1);  // missing axes input and noop. should be identity
+  test.AddInput<float>("data", {1, 2, 2},
+                       {1.0f, 2.0f,
+                        3.0f, 4.0f});
+  test.AddOptionalInputEdge<int64_t>();  // missing optional axes input
+  test.AddOutput<float>("reduced", {1, 2, 2},
+                        {1.0f, 2.0f,
+                         3.0f, 4.0f});
+
+  // TODO: OpenVINO doesn't support "axes" input in opset 13
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider});
+}
+
 TEST(ReductionOpTest, ReduceSum0DTensor) {
   OpTester test("ReduceSum");
   test.AddInput<float>("data", {}, {2});
@@ -5940,5 +5956,18 @@ TEST(ReductionOpTest, empty_set_ReduceSumSquare) {
 TEST(ReductionOpTest, empty_set_ReduceSumSquare_13) {
   test_empty_set("ReduceSumSquare", 13, false, 0.0f);
 }
+
+TEST(ReductionOpTest, MissingOptionalAxes) {
+  OpTester test("ReduceMax", 18);
+  test.AddInput<float>("data", {2, 2},
+                       {1.0f, 4.0f,
+                        3.0f, 2.0f});
+  test.AddOptionalInputEdge<int64_t>();
+  test.AddOutput<float>("reduced", {1, 1}, {4.0f});
+
+  // OpenVINO doesn't support "axes" input
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider});
+}
+
 }  // namespace test
 }  // namespace onnxruntime