license import { ArgMin , ArgMinAttrs , ArgMinInputs , KernelConfig , scalar , Tensor } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const argMinConfig = { kernelName : ArgMin , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { x } = args . inputs as ArgMinInputs ; const backend = args . backend as NodeJSKernelBackend ; const { axis } = args . attrs as { } as ArgMinAttrs ; const toDispose = [ ] ; let xInput = x ; if ( x . dtype === 'bool' ) { xInput = ( x as Tensor ) . toInt ( ) ; toDispose . push ( xInput as Tensor ) ; } const axisScalar = scalar ( axis , 'int32' ) ; toDispose . push ( axisScalar ) ; const opAttrs = [ createTensorsTypeOpAttr ( 'T' , xInput . dtype ) , createTensorsTypeOpAttr ( 'Tidx' , 'int32' ) , createTensorsTypeOpAttr ( 'output_type' , 'int32' ) ] ; const res = backend . executeSingleOutput ( ArgMin , opAttrs , [ xInput , axisScalar ] ) ; toDispose . forEach ( t => t . dispose ( ) ) ; return res ; } } ;