license import { ENGINE } from '../engine' ; import { GatherV2 , GatherV2Attrs , GatherV2Inputs } from '../kernel_names' ; import { NamedAttrMap } from '../kernel_registry' ; import { Tensor } from '../tensor' ; import { NamedTensorMap } from '../tensor_types' ; import { convertToTensor } from '../tensor_util_env' ; import { TensorLike } from '../types' ; import { op } from './operation' ; param x param indices param axis param batchDims doc function gather_ < T extends Tensor > ( x , indices , axis = 0 , batchDims = 0 ) { const $x = convertToTensor ( x , 'x' , 'gather' ) ; const $indices = convertToTensor ( indices , 'indices' , 'gather' , 'int32' ) ; const inputs = { x : $x , indices : $indices } ; const attrs = { axis , batchDims } ; return ENGINE . runKernel ( GatherV2 , inputs as { } as NamedTensorMap , attrs as { } as NamedAttrMap ) ; } export const gather = op ( { gather_ } ) ;