check 0 dimension

microsoft · Oct 24, 2024 · 8c35618 · 8c35618
1 parent 08cc261
commit 8c35618
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/onnxruntime/core/providers/webnn/builders/helper.cc b/onnxruntime/core/providers/webnn/builders/helper.cc
@@ -91,6 +91,10 @@ bool IsInputSupported(const NodeArg& input, const std::string& parent_name, cons
                             << input_name;
       return false;
     }
+    if (dim.dim_value() == 0) {
+      LOGS(logger, VERBOSE) << "The shape of [" << node_arg_name << "] has 0 dimension which is not supported";
+      return false;
+    }
   }
 
   return true;

diff --git a/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/expand_op_builder.cc
@@ -88,6 +88,10 @@ bool ExpandOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers
     LOGS(logger, VERBOSE) << "Cannot get shape.";
     return false;
   }
+  if (std::any_of(new_shape.begin(), new_shape.end(), [](int64_t dimension) { return dimension == 0; })) {
+    LOGS(logger, VERBOSE) << "Expand does not support new shape with 0 dimension.";
+    return false;
+  }
 
   std::vector<int64_t> input_shape;
   if (!GetShape(*input_defs[0], input_shape, logger)) {