license import { buffer , DataType , Rank , TensorBuffer } from '@tensorflow/tfjs-core' ; export function gatherV2Impl < R extends Rank , D extends DataType > ( xBuf , indicesBuf , flattenOutputShape ) { const outBuf = buffer ( flattenOutputShape , xBuf . dtype ) ; for ( let i = 0 ; i < outBuf . size ; ++ i ) { const newLoc = outBuf . indexToLoc ( i ) ; const originalLoc = newLoc . slice ( ) ; const batchIdx = originalLoc [ 0 ] ; const indicesIdx = originalLoc [ 2 ] ; const indicesIndex = indicesBuf . locToIndex ( [ batchIdx , indicesIdx ] ) ; originalLoc [ 2 ] = indicesBuf . values [ indicesIndex ] as number ; const originalIndex = xBuf . locToIndex ( originalLoc ) ; if ( 0 <= originalIndex && originalIndex < xBuf . values . length ) { outBuf . values [ i ] = xBuf . values [ originalIndex ] ; } } return outBuf as TensorBuffer < R , D > ; }