Function topKImpl

tfjs-backend-cpu/src/kernels/TopK_impl.ts:96–140

(
    x: TypedArray, xShape: number[], xDtype: NumericDataType, k: number,
    sorted: boolean)

}

export function topKImpl<T extends Tensor, R extends Rank>(
    x: TypedArray, xShape: number[], xDtype: NumericDataType, k: number,
    sorted: boolean):
    [TensorBuffer<R, NumericDataType>, TensorBuffer<R, 'int32'>] {
  // Reshape into a 2d tensor [batch, lastDim] and compute topk along lastDim.
  const lastDim = xShape[xShape.length - 1];
  const [batch, size] = [x.length / lastDim, lastDim];
  const allTopKVals = util.getTypedArrayFromDType(xDtype, batch * k);
  const allTopKIndices = util.getTypedArrayFromDType('int32', batch * k);

  for (let b = 0; b < batch; b++) {
    const offset = b * size;
    const vals = x.subarray(offset, offset + size);

    let valAndInd: Pair[] = new Array(vals.length);
    vals.forEach(
        (value: number, index: number) => valAndInd[index] = {value, index});

    if (k < valAndInd.length) {
      select(valAndInd, k);
      valAndInd = valAndInd.slice(0, k);
    }

    if (sorted) {
      valAndInd.sort(comparePair);
    }
    
    const outOffset = b * k;
    const topKVals = allTopKVals.subarray(outOffset, outOffset + k);
    const topKIndices = allTopKIndices.subarray(outOffset, outOffset + k);
    for (let i = 0; i < k; i++) {
      topKVals[i] = valAndInd[i].value;
      topKIndices[i] = valAndInd[i].index;
    }
  }
  // Reshape back to the original input shape, except that the last
  // dimension is k.
  const outputShape = xShape.slice();
  outputShape[outputShape.length - 1] = k;

  return [
    buffer(outputShape as ShapeMap[R], xDtype, allTopKVals),
    buffer(outputShape as ShapeMap[R], 'int32', allTopKIndices)
  ];
}