Method constructor

tfjs-backend-webgl/src/gather_gpu.ts:29–43

(aShape: GatherShape, outputShape: GatherShape)

  rank: number;

  constructor(aShape: GatherShape, outputShape: GatherShape) {
    this.outputShape = outputShape;
    this.rank = outputShape.length;
    const dtype = getCoordsDataType(this.rank);
    const sourceCoords = getSourceCoords(aShape, 2);

    this.userCode = `
      void main() {
        ${dtype} resRC = getOutputCoords();
        int index = int(getIndices(resRC.x, resRC.z));
        float inBounds = (index >= 0) && (index < ${aShape[2]}) ? 1.0 : 0.0;
        setOutput(inBounds * getA(${sourceCoords}));
      }
    `;
  }
}

// The input and output are always flattened into rank 4 tensors.