From 66df11769c127565dcd17a715b2e08b3bffd9eee Mon Sep 17 00:00:00 2001
From: Hariharan Seshadri <shariharan91@gmail.com>
Date: Wed, 16 Aug 2023 11:24:26 -0700
Subject: [PATCH] [JS/WebGPU] Expand operator fixes (#17137)

---
 js/web/lib/wasm/jsep/webgpu/ops/expand.ts |  6 +-
 js/web/test/data/ops/expand.jsonc         | 89 +++++++++++++++++++++++
 js/web/test/suite-test-list.jsonc         |  1 +
 3 files changed, 93 insertions(+), 3 deletions(-)
 create mode 100644 js/web/test/data/ops/expand.jsonc

diff --git a/js/web/lib/wasm/jsep/webgpu/ops/expand.ts b/js/web/lib/wasm/jsep/webgpu/ops/expand.ts
index 04809e78c1..b07fe3a90f 100644
--- a/js/web/lib/wasm/jsep/webgpu/ops/expand.ts
+++ b/js/web/lib/wasm/jsep/webgpu/ops/expand.ts
@@ -65,8 +65,7 @@ const createExpandProgramInfo = (metadata: ProgramMetadata, inputs: readonly Ten
     let outputIndices = ${output.offsetToIndices('global_idx')};
     var inputIndices: ${input.type.indices};
     for (var i = 0; i < ${inputShape.length}; i++) {
-      if (inputShape[i] == 1) {
-        // TODO: IndicesHelper should offer uniform way to get/set indices for all ranks
+      if (${input.indicesGet('inputShape', 'i')} == 1) {
         ${input.indicesSet('inputIndices', 'i', 0)}
       } else {
         ${
@@ -86,7 +85,8 @@ const createExpandProgramInfo = (metadata: ProgramMetadata, inputs: readonly Ten
 
 export const expand = (context: ComputeContext): void => {
   validateInputs(context.inputs);
-  const cacheHint = context.inputs.map(x => x.dims.toString()).join('_');
+  const outputShape = Array.from(context.inputs[1].getBigInt64Array(), Number);
+  const cacheHint = outputShape.toString();
   context.compute(
       {...expandProgramMetadata, cacheHint, get: () => createExpandProgramInfo(expandProgramMetadata, context.inputs)},
       {inputs: [0]});
diff --git a/js/web/test/data/ops/expand.jsonc b/js/web/test/data/ops/expand.jsonc
new file mode 100644
index 0000000000..460122b4e0
--- /dev/null
+++ b/js/web/test/data/ops/expand.jsonc
@@ -0,0 +1,89 @@
+[
+  {
+    "name": "Expand 1 - float32",
+    "operator": "Expand",
+    "attributes": [],
+    "cases": [
+      {
+        "name": "Expand 1 - float32",
+        "inputs": [
+          {
+            "data": [1],
+            "dims": [1, 1],
+            "type": "float32"
+          },
+          {
+            "data": [1, 4],
+            "dims": [2],
+            "type": "int64"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [1, 1, 1, 1],
+            "dims": [1, 4],
+            "type": "float32"
+          }
+        ]
+      }
+    ]
+  },
+  {
+    "name": "Expand 2 - float32",
+    "operator": "Expand",
+    "attributes": [],
+    "cases": [
+      {
+        "name": "Expand 2 - float32",
+        "inputs": [
+          {
+            "data": [1],
+            "dims": [1, 1],
+            "type": "float32"
+          },
+          {
+            "data": [1, 6],
+            "dims": [2],
+            "type": "int64"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [1, 1, 1, 1, 1, 1],
+            "dims": [1, 6],
+            "type": "float32"
+          }
+        ]
+      }
+    ]
+  },
+  {
+    "name": "Expand 1D - float32",
+    "operator": "Expand",
+    "attributes": [],
+    "cases": [
+      {
+        "name": "Expand 1D - float32",
+        "inputs": [
+          {
+            "data": [1],
+            "dims": [1],
+            "type": "float32"
+          },
+          {
+            "data": [4],
+            "dims": [1],
+            "type": "int64"
+          }
+        ],
+        "outputs": [
+          {
+            "data": [1, 1, 1, 1],
+            "dims": [4],
+            "type": "float32"
+          }
+        ]
+      }
+    ]
+  }
+]
diff --git a/js/web/test/suite-test-list.jsonc b/js/web/test/suite-test-list.jsonc
index 72300ebb3e..e6224bd3df 100644
--- a/js/web/test/suite-test-list.jsonc
+++ b/js/web/test/suite-test-list.jsonc
@@ -1334,6 +1334,7 @@
       //"depth-to-space.jsonc",
       //"equal.jsonc",
       "exp.jsonc",
+      "expand.jsonc",
       "floor.jsonc",
       "gemm.jsonc",
       "global-average-pool.jsonc",