Function reduce

tfjs-backend-webgpu/src/kernel_utils/reduce.ts:35–103

(
    x: TensorInfo, axis: number|number[], keepDims: boolean,
    reduceType: ReduceTypes, backend: WebGPUBackend)

};

export function reduce(
    x: TensorInfo, axis: number|number[], keepDims: boolean,
    reduceType: ReduceTypes, backend: WebGPUBackend): TensorInfo {
  const xRank = x.shape.length;
  const toDispose = [];

  const origAxes = util.parseAxisParam(axis, x.shape);
  let axes = origAxes;
  const permutedAxes = backend_util.getAxesPermutation(axes, xRank);

  let input = x;
  if (permutedAxes != null) {
    input = transpose({inputs: {x}, attrs: {perm: permutedAxes}, backend});
    axes = backend_util.getInnerMostAxes(axes.length, xRank);
    toDispose.push(input);
  }

  backend_util.assertAxesAreInnerMostDims(reduceType, axes, xRank);

  const [reduceOutShape, reduceShape] =
      backend_util.computeOutAndReduceShapes(input.shape, axes);
  let resOutShape = reduceOutShape;
  if (keepDims) {
    // rather than reshape at the end, set the target shape here.
    resOutShape = backend_util.expandShapeToKeepDim(reduceOutShape, origAxes);
  }

  let res;
  if ((reduceType === 'max' || reduceType === 'prod') &&
      backend.shouldExecuteOnCPU([input])) {
    const xVals = backend.tensorMap.get(input.dataId).values as TypedArray;
    switch (reduceType) {
      case 'max':
        const outValues = maxImplCPU(
            xVals, util.sizeFromShape(reduceShape), resOutShape, x.dtype);
        res = backend.makeTensorInfo(resOutShape, x.dtype, outValues);
        break;
      case 'prod':
        const {outVals, outShape, outDtype} =
            prodImplCPU(input.shape, input.dtype, xVals, axes);
        res = backend.makeTensorInfo(outShape, outDtype, outVals);
        break;
      default:
        throw new Error(
            `${reduceType} CPU implementation is not yet supported.`);
    }
  } else {
    const inSize = util.sizeFromShape(reduceShape);
    const xSize = util.sizeFromShape(input.shape);
    const batchSize = xSize / inSize;

    const reduceInfo = {windowSize: inSize, inSize, batchSize, outSize: 1};
    const dtype = RETURN_TYPES[reduceType] || sumOutType(x.dtype);
    const uniformData = [
      {type: 'int32', data: [inSize]},
    ];
    const program = new ReduceProgram(
        reduceInfo, reduceType, backend.device.limits.maxComputeWorkgroupSizeX);