license import { backend_util , GatherNd , GatherNdInputs , KernelConfig , KernelFunc , TensorInfo , TypedArray , util } from '@tensorflow/tfjs-core' ; import { MathBackendCPU } from '../backend_cpu' ; import { gatherNdImpl } from './GatherNd_Impl' ; export function gatherNd ( args ) { const { inputs , backend } = args ; const { params , indices } = inputs ; const paramsSize = util . sizeFromShape ( params . shape ) ; const indicesShape = indices . shape ; const sliceRank = indicesShape [ indicesShape . length - 1 ] ; const [ resultShape , numSlices , sliceSize , strides ] = backend_util . prepareAndValidate ( params , indices ) ; if ( numSlices === 0 ) { return backend . makeTensorInfo ( resultShape , params . dtype , [ ] ) ; } const indicesData = backend . data . get ( indices . dataId ) . values as TypedArray ; const paramsBuf = backend . bufferSync ( params ) ; const outBuf = gatherNdImpl ( indicesData , paramsBuf , params . dtype , numSlices , sliceRank , sliceSize , strides , params . shape , paramsSize ) ; return backend . makeTensorInfo ( resultShape , params . dtype , outBuf . values ) ; } export const gatherNdConfig = { kernelName : GatherNd , backendName : 'cpu' , kernelFunc : gatherNd as { } as KernelFunc } ;