From ede50ebb3268ae9eb9ece413453107739282fd8e Mon Sep 17 00:00:00 2001
From: Shiyi Zou <shiyi.zou@intel.com>
Date: Wed, 23 Oct 2024 15:17:40 +0800
Subject: [PATCH] check 0 dimension

---
 onnxruntime/core/providers/webnn/builders/helper.cc           | 4 ++++
 .../core/providers/webnn/builders/impl/expand_op_builder.cc   | 4 ++++
 .../core/providers/webnn/builders/impl/reshape_op_builder.cc  | 2 +-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/webnn/builders/helper.cc b/onnxruntime/core/providers/webnn/builders/helper.cc
index b90c7d76a6507..e21d73f19e43d 100644
--- 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
index c8cea833983b1..e74ad4dd30627 100644
--- 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
index a7911683f0355..453c4c6e5cd5d 100644
--- 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;
   }