[WebGPU] Fix shader key for ScatterProgram (tensorflow#7932)

* Add uniforms * Update webgpu_program.ts * Update webgpu_program.ts * Revert "Add uniforms" This reverts commit 645802f. * Revert "Update webgpu_program.ts" This reverts commit 58ea96f. * Revert "Update webgpu_program.ts" This reverts commit 32386ac. * Add key to scatter webgpu program
fengwuyao · Sep 25, 2023 · 73b2fd1 · 73b2fd1
1 parent a7bec12
commit 73b2fd1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tfjs-backend-webgpu/src/scatter_webgpu.ts b/tfjs-backend-webgpu/src/scatter_webgpu.ts
@@ -48,8 +48,9 @@ export class ScatterProgram implements WebGPUProgram {
     this.dispatch =
         computeDispatch(this.dispatchLayout, flattenXShape, this.workgroupSize);
     this.sliceDimGreaterThanOne = sliceDim > 1;
-    this.shaderKey = `scatter_${indicesRank}_${updatesRank}_${
-        this.sliceDimGreaterThanOne}_${outputDtype}_${sumDupeIndices}`;
+    this.shaderKey =
+        `scatter_${indicesRank}_${updatesRank}_${this.sliceDimGreaterThanOne}_${
+            outputDtype}_${sumDupeIndices}_${strides.length}`;
     const stridesType = getCoordsDataType(strides.length);
     this.uniforms =
         `sliceDim : i32, strides: ${stridesType}, updatesSize: i32,`;