[js/webgpu] Fix wrong shape in maxpool

Bug: microsoft#21386
axinging · Jul 30, 2024 · 64f1849 · 64f1849
1 parent 5d78b9a
commit 64f1849
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/js/web/lib/wasm/jsep/webgpu/ops/pool.ts b/js/web/lib/wasm/jsep/webgpu/ops/pool.ts
@@ -31,8 +31,14 @@ const getAdjustedPoolAttributesAndOutputShape = <AttributeType extends AveragePo
     inputShapeAsChannelFirst.splice(1, 0, inputShapeAsChannelFirst.pop()!);  // Move channel to the second position.
   }
   const hasDilations = Object.hasOwnProperty.call(attributes, 'dilations');
-  const kernelShape = attributes.kernelShape.slice();
-  const strides = attributes.strides.slice();
+  let kernelShape = attributes.kernelShape.slice();
+  if (kernelShape[kernelShape.length - 1] === 0) {
+    kernelShape.pop();
+  }
+  let strides = attributes.strides.slice();
+  if (strides[strides.length - 1] === 0) {
+    strides.pop();
+  }
   const dilations: number[] = hasDilations ? (attributes as MaxPoolAttributes).dilations.slice() : [];
   const pads = attributes.pads.slice();
   PoolConvUtil.adjustPoolAttributes(isGlobalOperator, inputShapeAsChannelFirst, kernelShape, strides, dilations, pads);

diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc
@@ -1370,6 +1370,7 @@
       "matmul.jsonc",
       "matmulnbits.jsonc",
       "matmul-broadcast.jsonc",
+      "max-pool.jsonc",
       "mul.jsonc",
       "mul_int32.jsonc",
       "multihead-attention.jsonc",