Function transpose

tfjs-backend-cpu/src/kernels/Transpose.ts:25–48

(args: {
  inputs: TransposeInputs,
  attrs: TransposeAttrs,
  backend: MathBackendCPU
})

import {transposeImpl} from './Transpose_impl';

export function transpose(args: {
  inputs: TransposeInputs,
  attrs: TransposeAttrs,
  backend: MathBackendCPU
}): TensorInfo {
  const {inputs, attrs, backend} = args;
  const {x} = inputs;
  const {perm} = attrs;

  assertNotComplex(x, 'transpose');

  const xRank = x.shape.length;

  const newShape: number[] = new Array(xRank);
  for (let i = 0; i < newShape.length; i++) {
    newShape[i] = x.shape[perm[i]];
  }

  const values = backend.data.get(x.dataId).values as TypedArray;
  const result = transposeImpl(values, x.shape, x.dtype, perm, newShape);

  const dataId = backend.write(result, newShape, x.dtype);
  return {dataId, shape: newShape, dtype: x.dtype};
}

export const transposeConfig: KernelConfig = {
  kernelName: Transpose,