Function transpose

tfjs-backend-webgpu/src/kernels/Transpose.ts:26–53

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

import {TransposeProgram} from '../transpose_webgpu';

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

  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]];
  }
  if (backend.shouldExecuteOnCPU([x])) {
    const xData = webgpuBackend.tensorMap.get(x.dataId);
    const values = xData.values as TypedArray;
    const outValues = cpuTranspose(values, x.shape, x.dtype, perm, newShape);
    return backend.makeTensorInfo(newShape, x.dtype, outValues);
  }
  if (x.shape.length === 2 && util.arraysEqual(perm, [1, 0])) {
    const program = new TransposeSharedProgram(x.shape, perm);
    return webgpuBackend.runWebGPUProgram(program, [x], x.dtype);
  }
  const program = new TransposeProgram(x.shape, perm);
  return webgpuBackend.runWebGPUProgram(program, [x], x.dtype);
}

export const transposeConfig: KernelConfig = {
  kernelName: Transpose,