license import { KernelConfig , KernelFunc , TensorInfo , Transpose , TransposeAttrs , TransposeInputs , TypedArray } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { transposeImpl } from './Transpose_impl' ; import { transposeImplCPU as cpuTranspose } from './Transpose_impl' ; export function transpose ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { perm } = attrs ; const webglBackend = backend ; const xRank = x . shape . length ; const newShape = new Array ( xRank ) ; for ( let i = 0 ; i < newShape . length ; i ++ ) { newShape [ i ] = x . shape [ perm [ i ] ] ; } let out ; if ( webglBackend . shouldExecuteOnCPU ( [ x ] ) ) { const xTexData = webglBackend . texData . get ( x . dataId ) ; const values = xTexData . values as TypedArray ; const outValues = cpuTranspose ( values , x . shape , x . dtype , perm , newShape ) ; out = webglBackend . makeTensorInfo ( newShape , x . dtype ) ; const outData = webglBackend . texData . get ( out . dataId ) ; outData . values = outValues ; } else { out = transposeImpl ( x , perm , webglBackend ) ; } return out ; } export const transposeConfig = { kernelName : Transpose , backendName : 'webgl' , kernelFunc : transpose as { } as KernelFunc } ;