license import { backend_util , KernelConfig , scalar , SplitV , SplitVAttrs , SplitVInputs , Tensor , tensor1d , tidy , util } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const splitVConfig = { kernelName : SplitV , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { x } = args . inputs as SplitVInputs ; const backend = args . backend as NodeJSKernelBackend ; const { numOrSizeSplits , axis } = args . attrs as { } as SplitVAttrs ; const $axis = util . parseAxisParam ( axis , x . shape ) [ 0 ] ; const splitSizes = backend_util . prepareSplitSize ( x , numOrSizeSplits , $axis ) ; const opAttrs = [ { name : 'num_split' , type : backend . binding . TF_ATTR_INT , value : splitSizes . length } , createTensorsTypeOpAttr ( 'T' , x as Tensor ) , { name : 'Tlen' , type : backend . binding . TF_ATTR_TYPE , value : backend . binding . TF_INT32 } ] ; const inputs = [ x ] ; return tidy ( ( ) => { inputs . push ( tensor1d ( splitSizes , 'int32' ) ) ; inputs . push ( scalar ( $axis , 'int32' ) ) ; return backend . executeMultipleOutputs ( SplitV , opAttrs , inputs , splitSizes . length ) ; } ) ; } } ;