Fixed bug in Flatten's axis (microsoft#18645)

Flatten's axis is in the range [-r, r] rather than [-r, r-1].
axinging · Dec 1, 2023 · 73a2eb8 · 73a2eb8
1 parent 6781b6c
commit 73a2eb8
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/onnxruntime/core/providers/webnn/builders/impl/flatten_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/flatten_op_builder.cc
@@ -36,7 +36,11 @@ Status FlattenOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
   int64_t rank = input_shape.size();
   NodeAttrHelper helper(node);
   int64_t axis = helper.Get("axis", 1);
-  axis = HandleNegativeAxis(axis, rank);
+  ORT_ENFORCE(axis >= -rank && axis <= rank, "axis ", axis,
+              " is not in valid range [-", rank, ",", rank, "]");
+  if (axis < 0) {
+    axis += rank;
+  }
 
   // Use WebNN's reshape to implement Flatten.
   int64_t num_pre_axis_elements = std::accumulate(