webgpu: Fix a bug in softmax (#7607)

tensorflow · Apr 20, 2023 · 78a3a03 · 78a3a03
1 parent f35881e
commit 78a3a03
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/tfjs-backend-webgpu/src/softmax_webgpu.ts b/tfjs-backend-webgpu/src/softmax_webgpu.ts
@@ -59,16 +59,17 @@ export class SoftmaxProgram implements WebGPUProgram {
       }
       workgroupBarrier();
 
-      let reduceSize = min(cols, blockSize);
-      for (var currSize = reduceSize >> 1;  currSize > 0; currSize = currSize >> 1) {
+      var reduceSize = min(cols, blockSize);
+      for (var currSize = reduceSize >> 1;  currSize > 0; currSize = reduceSize >> 1) {
+        reduceSize = currSize + (reduceSize & 1);
         if (tid < currSize) {
-          buf[tid] = max(buf[tid], buf[tid + currSize]);
+          buf[tid] = max(buf[tid], buf[tid + reduceSize]);
         }
         workgroupBarrier();
       }
 
       if (tid == 0) {
-        rowMaxShared = max(buf[0], buf[reduceSize - 1]);
+        rowMaxShared = buf[0];
       }
       workgroupBarrier();
 

diff --git a/tfjs-core/src/ops/softmax_test.ts b/tfjs-core/src/ops/softmax_test.ts
@@ -39,6 +39,12 @@ describeWithFlags('softmax', ALL_ENVS, () => {
     expectArraysClose(await y.data(), [0.5, 0.5]);
   });
 
+  it('odd number of inputs', async () => {
+    const y = tf.softmax(tf.tensor1d([-400, -400, 0, -400, -400, -400, -400]));
+
+    expectArraysClose(await y.data(), [0, 0, 1, 0, 0, 0, 0]);
+  });
+
   it('Huge difference between probabilities', async () => {
     const y = tf.softmax(tf.tensor1d([-1000, +1000]));