license /** An implementation of the TopK kernel shared between webgl and cpu. */ import { buffer , NumericDataType , Rank , ShapeMap , Tensor , TensorBuffer , TypedArray , util } from '@tensorflow/tfjs-core' ; type Pair = { value : number ; index : number ; } ; const comparePair = ( a , b ) => { const valueDiff = b . value - a . value ; return valueDiff === 0 ? a . index - b . index : valueDiff ; } ; param array param left param right param k function select ( array , k , left = 0 , right = array . length - 1 ) { while ( right > left ) { if ( right - left > 600 ) { const n = right - left + 1 ; const i = k - left + 1 ; const z = Math . log ( n ) ; const s = 0.5 * Math . exp ( 2 * z / 3 ) ; const sd = 0.5 * Math . sqrt ( z * s * ( n - s ) / n ) * Math . sign ( i - n / 2 ) ; const newLeft = Math . max ( left , Math . floor ( k - i * s / n + sd ) ) ; const newRight = Math . min ( right , Math . floor ( k + ( n - i ) * s / n + sd ) ) ; select ( array , k , newLeft , newRight ) ; } const t = array [ k ] ; let i = left ; let j = right ; util . swap ( array , left , k ) ; if ( comparePair ( array [ right ] , t ) > 0 ) { util . swap ( array , left , right ) ; } while ( i < j ) { util . swap ( array , i , j ) ; i ++ ; j -- ; while ( comparePair ( array [ i ] , t ) < 0 ) { i = i + 1 ; } while ( comparePair ( array [ j ] , t ) > 0 ) { j = j - 1 ; } } if ( comparePair ( array [ left ] , t ) === 0 ) { util . swap ( array , left , j ) ; } else { j = j + 1 ; util . swap ( array , j , right ) ; } if ( j <= k ) { left = j + 1 ; } if ( k <= j ) { right = j - 1 ; } } } export function topKImpl < T extends Tensor , R extends Rank > ( x , xShape , xDtype , k , sorted ) { const lastDim = xShape [ xShape . length - 1 ] ; const [ batch , size ] = [ x . length / lastDim , lastDim ] ; const allTopKVals = util . getTypedArrayFromDType ( xDtype , batch * k ) ; const allTopKIndices = util . getTypedArrayFromDType ( 'int32' , batch * k ) ; for ( let b = 0 ; b < batch ; b ++ ) { const offset = b * size ; const vals = x . subarray ( offset , offset + size ) ; let valAndInd = new Array ( vals . length ) ; vals . forEach ( ( value , index ) => valAndInd [ index ] = { value , index } ) ; if ( k < valAndInd . length ) { select ( valAndInd , k ) ; valAndInd = valAndInd . slice ( 0 , k ) ; } if ( sorted ) { valAndInd . sort ( comparePair ) ; } const outOffset = b * k ; const topKVals = allTopKVals . subarray ( outOffset , outOffset + k ) ; const topKIndices = allTopKIndices . subarray ( outOffset , outOffset + k ) ; for ( let i = 0 ; i < k ; i ++ ) { topKVals [ i ] = valAndInd [ i ] . value ; topKIndices [ i ] = valAndInd [ i ] . index ; } } const outputShape = xShape . slice ( ) ; outputShape [ outputShape . length - 1 ] = k ; return [ buffer ( outputShape as ShapeMap [ R ] , xDtype , allTopKVals ) , buffer ( outputShape as ShapeMap [ R ] , 'int32' , allTopKIndices ) ] ; }