check 0 dimension

shiyi9801 · Oct 24, 2024 · ede50eb · ede50eb
1 parent 08cc261
commit ede50eb
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 1 deletion.
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)) {

diff --git a/onnxruntime/core/providers/webnn/builders/impl/reshape_op_builder.cc b/onnxruntime/core/providers/webnn/builders/impl/reshape_op_builder.cc
@@ -92,7 +92,7 @@ bool ReshapeOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializer
 
   const int64_t* raw_new_shape = reinterpret_cast<const int64_t*>(unpacked_tensor.data());
   const auto& perm_dims = perm_tensor.dims();
-  if (perm_dims.empty() || perm_dims[0] == 0) {
+  if (perm_dims.empty()) {
     LOGS(logger, VERBOSE) << "New shape of reshape cannot be empty";
     return false;
   }