Method constructor

tfjs-backend-webgpu/src/argminmax_webgpu.ts:36–66

(inputShape: number[], axis: number, reduceType: 'min'|'max')

  private type: string;

  constructor(inputShape: number[], axis: number, reduceType: 'min'|'max') {
    const axes = [axis];

    this.op = reduceType === 'min' ? '<' : '>';

    // |outShape| is the shape with the removed axis
    const [outputShape, reduceShape] =
        backend_util.computeOutAndReduceShapes(inputShape, axes);

    this.outputShape = outputShape.length === 0 ? [1] : outputShape;
    this.dispatchLayout = flatDispatchLayout(this.outputShape);
    // The shared algorithm is mainly used for large reduce size. It fully
    // utilizes the threads in one workgroup to do the reduction. However,
    // when the reduce size is very small, it's better to use the plain
    // algorithm to reduce the number of workgroups to speedup. The threthold
    // can be further tuned.
    if (util.sizeFromShape(reduceShape) < 32) {
      this.type = 'plain';
      this.dispatch = computeDispatch(
          this.dispatchLayout, this.outputShape, this.workgroupSize);
    } else {
      this.type = 'shared';
      // A work group only outputs a data, so we transfer [1, 1, 1] to compute
      // dispatch size.
      this.dispatch =
          computeDispatch(this.dispatchLayout, this.outputShape, [1, 1, 1]);
    }

    this.inputShape = inputShape;
    this.shaderKey = `argMinMax_${this.op}_${this.type}`;
  }

  getUserCode(): string {
    const workgroupSizeX = this.workgroupSize[0];