Method constructor

tfjs-backend-webgl/src/cum_gpu.ts:30–65

(
      public op: CumOpType, public outputShape: number[], exclusive: boolean,
      reverse: boolean)

  customUniforms = [{name: 'index', type: 'float' as UniformType}];

  constructor(
      public op: CumOpType, public outputShape: number[], exclusive: boolean,
      reverse: boolean) {
    const rank = this.outputShape.length;
    const initVal = this.op === CumOpType.Prod ? '1.0' : '0.0';
    const val =
        exclusive ? initVal : `getX(${getCoords(rank, 'coords', this.op)})`;
    const length = this.outputShape[this.outputShape.length - 1];
    let condition = '';
    let idxString = '';
    // When exclusive is set, the cum op becomes roll op that copies the
    // value from the previous index based on the direction specified by the
    // reverse flag.
    if (exclusive) {
      condition = reverse ? `end != ${length - 1}` : 'end != 0';
      idxString = reverse ? 'end + 1' : 'end - 1';
    } else {
      condition = reverse ? `end + pow2 < ${length}` : 'end >= pow2';
      idxString = (reverse ? 'end + pow2' : 'end - pow2');
    }

    this.userCode = `
      void main() {
        ${getCoordsDataType(rank)} coords = getOutputCoords();
        int end = ${getFinalCoord(rank, 'coords', this.op)};
        float val = ${val};
        int pow2 = int(pow(2.0, index));
        if (${condition}) {
          int idx = ${idxString};
          ${getFinalCoord(rank, 'coords', this.op)} = idx;
          val ${this.op}= getX(${getCoords(rank, 'coords', this.op)});
        }
        setOutput(val);
      }
    `;
  }
}

function getCoords(rank: number, name: string, op: CumOpType): string {