import { Tensor } from 'onnxruntime-web' ; import { NumberDataType } from './yoloPostprocess' ; import { arrayCopyHelper , ShapeUtil , TypedArrayUtil } from './yoloPostprocessUtils' ; export function transpose ( x , perm ? ) { const inputDims = x . dims ? x . dims : [ x . data . length ] ; const rank = inputDims . length ; let finalPerm = new Array < number > ( rank ) ; if ( perm && perm . length === rank ) { finalPerm = perm ; } else { for ( let i = 0 ; i < rank ; i ++ ) { finalPerm [ i ] = rank - i - 1 ; } } const outputDims = new Array < number > ( rank ) ; const stride = new Array < number > ( rank ) ; for ( let i = 0 ; i < rank ; i ++ ) { const inpDim = finalPerm [ i ] ; outputDims [ i ] = inputDims [ inpDim ] ; if ( inpDim + 1 < rank ) { stride [ i ] = ShapeUtil . sizeFromDimension ( inputDims , inpDim + 1 ) ; } else { stride [ i ] = 1 ; } } const output = new Tensor ( x . type , TypedArrayUtil . createTypedArray ( x . type , ShapeUtil . size ( outputDims ) ) , outputDims ) ; const X = x . data as NumberDataType ; const Y = output . data as NumberDataType ; let numAxesInPrefix = 0 ; let suffixBlocksize = 1 ; let prefixBlocksize = 1 ; let isSuffix = true ; for ( let i = rank - 1 ; i >= 0 ; -- i ) { const inpAxis = finalPerm [ i ] ; if ( isSuffix && ( inpAxis === i ) ) { suffixBlocksize *= inputDims [ inpAxis ] ; } else { isSuffix = false ; prefixBlocksize *= inputDims [ inpAxis ] ; ++ numAxesInPrefix ; } } if ( prefixBlocksize === 1 ) { doTransposeSingleBlock ( suffixBlocksize , Y , X ) ; } else if ( suffixBlocksize === 1 ) { doTransposeEltWise ( numAxesInPrefix , outputDims , prefixBlocksize , stride , Y , X ) ; } else { doTranspose ( numAxesInPrefix , outputDims , prefixBlocksize , suffixBlocksize , stride , Y , X ) ; } return output ; } function doTranspose ( numAxes , targetDims , numBlocks , numElementsInBlock , stride , target , source ) { const targetIndex = new Array < number > ( numAxes ) . fill ( 0 ) ; const startSourceIndex = 0 ; let startTargetIndex = 0 ; for ( let i = 0 ; i < numBlocks ; ++ i ) { const sizeOffset = ShapeUtil . computeOffset ( targetIndex , stride , numAxes ) ; arrayCopyHelper ( target , source , startTargetIndex , startSourceIndex + sizeOffset , numElementsInBlock ) ; ShapeUtil . incrementIndex ( targetIndex , targetDims , numAxes ) ; startTargetIndex += numElementsInBlock ; } } function doTransposeEltWise ( numAxes , targetDims , numBlocks , stride , target , source ) { const targetIndex = new Array < number > ( numAxes ) . fill ( 0 ) ; let startTargetIndex = 0 ; for ( let i = 0 ; i < numBlocks ; ++ i ) { const sourceOffset = ShapeUtil . computeOffset ( targetIndex , stride , numAxes ) ; target [ startTargetIndex ++ ] = source [ sourceOffset ] ; ShapeUtil . incrementIndex ( targetIndex , targetDims , numAxes ) ; } } function doTransposeSingleBlock ( numElementsInBlock , target , source ) { arrayCopyHelper ( target , source , 0 , 0 , numElementsInBlock ) ; }