Function transposeImpl

tfjs-backend-cpu/src/kernels/Transpose_impl.ts:21–45

(
    xVals: TypedArray, xShape: number[], dtype: DataType, perm: number[],
    newShape: number[])

import {util} from '@tensorflow/tfjs-core';

export function transposeImpl(
    xVals: TypedArray, xShape: number[], dtype: DataType, perm: number[],
    newShape: number[]): TypedArray {
  const xRank = xShape.length;
  const xSize = util.sizeFromShape(xShape);
  const xStrides = util.computeStrides(xShape);
  const newStrides = util.computeStrides(newShape);

  const result = util.getTypedArrayFromDType(
      dtype as NumericDataType, util.sizeFromShape(newShape));

  for (let i = 0; i < xSize; ++i) {
    const loc = util.indexToLoc(i, xRank, xStrides);

    // Permute location.
    const newLoc: number[] = new Array(loc.length);
    for (let i = 0; i < newLoc.length; i++) {
      newLoc[i] = loc[perm[i]];
    }

    const newIndex = util.locToIndex(newLoc, xRank, newStrides);
    result[newIndex] = xVals[i];
  }
  return result;
}