Webgpu gather zero

tfjs-backend-webgl/src/webgl_ops_test.ts

@@ -873,21 +873,6 @@ describeWithFlags('depthwiseConv2d packed', PACKED_ENVS, () => {
   });
 });
 
-describeWithFlags('gather', WEBGL_ENVS, () => {
-  it('fills with zero when index is out of bound', async () => {
-    const t = tf.tensor2d([1, 11, 2, 22], [2, 2]);
-    const tInt = tf.tensor2d([1, 11, 2, 22], [2, 2], 'int32');
-
-    const index = tf.tensor1d([0, 1, 100, -1, 2, -4], 'int32');
-    const res = tf.gather(t, index);
-    const resInt = tf.gather(tInt, index);
-
-    const expected = [1, 11, 2, 22, 0, 0, 0, 0, 0, 0, 0, 0];
-    expectArraysClose(await res.data(), expected);
-    expectArraysClose(await resInt.data(), expected);
-  });
-});
-
 describeWithFlags('gather debug', WEBGL_ENVS, () => {
   it('throws if index is out of bound if in debug mode', async () => {
     const t = tf.tensor2d([1, 11, 2, 22], [2, 2]);

tfjs-backend-webgpu/src/gather_webgpu.ts

@@ -41,12 +41,14 @@ export class GatherProgram implements WebGPUProgram {
   }
 
   getUserCode(): string {
-    const sourceCoords = getSourceCoords(this.aShape, 'i32');
+    const sourceCoords = getSourceCoords(this.aShape);
     const userCode = `
       ${getMainHeaderAndGlobalIndexString()}
         if (index < uniforms.size) {
           let resRC = getCoordsFromIndex(index);
-          setOutputAtIndex(index, getA(${sourceCoords}));
+          let indexZ = i32(getIndices(resRC.x, resRC.z));
+          let inBounds = select(0.0, 1.0, indexZ >= 0 && indexZ < uniforms.aShape[2]);
+          setOutputAtIndex(index, inBounds * getA(${sourceCoords}));
         }
       }
     `;
@@ -55,12 +57,12 @@ export class GatherProgram implements WebGPUProgram {
 }
 
 // The input and output are always flattened into rank 4 tensors.
-function getSourceCoords(aShape: number[], typePrefix = 'int'): string {
+function getSourceCoords(aShape: number[]): string {
   const currentCoords = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w'];
   const sourceCoords = [];
   for (let i = 0; i < aShape.length; i++) {
     if (i === 2) {
-      sourceCoords.push(`${typePrefix}(getIndices(resRC.x, resRC.z))`);
+      sourceCoords.push('indexZ');
     } else {
       sourceCoords.push(`${currentCoords[i]}`);
     }

tfjs-backend-webgpu/src/kernels/GatherV2.ts

@@ -31,7 +31,8 @@ export function gatherV2(
   const {x, indices} = inputs;
   const {axis, batchDims} = attrs;
 
-  // Throw error when any index is out of bound.
+  // Unlike WebGL, WebGPU won't check if index is out of bound by calling
+  // backend.readSync() function in debug mode.
   const parsedAxis = util.parseAxisParam(axis, x.shape)[0];
 
   const shapeInfo = backend_util.segment_util.collectGatherOpShapeInfo(

tfjs-core/src/ops/gather_test.ts

@@ -19,7 +19,7 @@ import * as tf from '../index';
 import {ALL_ENVS, describeWithFlags} from '../jasmine_util';
 import {expectArraysClose} from '../test_util';
 
-describeWithFlags('gather', ALL_ENVS, () => {
+describeWithFlags('gather', ALL_ENVS, (env) => {
   it('1D (gather), scalar indices', async () => {
     const t = tf.tensor1d([1, 2, 3]);
 
@@ -579,4 +579,19 @@ describeWithFlags('gather', ALL_ENVS, () => {
     expect(numTensorsAfter).toBe(numTensorsBefore);
     expect(numDataIdAfter).toBe(numDataIdBefore);
   });
+
+  it('fills with zero when index is out of bound', async () => {
+    if (env.backendName === 'webgl' || env.backendName === 'webgpu') {
+      const t = tf.tensor2d([1, 11, 2, 22], [2, 2]);
+      const tInt = tf.tensor2d([1, 11, 2, 22], [2, 2], 'int32');
+
+      const index = tf.tensor1d([0, 1, 100, -1, 2, -4], 'int32');
+      const res = tf.gather(t, index);
+      const resInt = tf.gather(tInt, index);
+
+      const expected = [1, 11, 2, 22, 0, 0, 0, 0, 0, 0, 0, 0];
+      expectArraysClose(await res.data(), expected);
+      expectArraysClose(await resInt.data(), expected);
+    }
+  });
 });