[JS/WebGPU] Expand operator fixes (#17137)

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]});

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"
+          }
+        ]
+      }
+    ]
+  }
+]

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",