license import { KernelConfig , Range , RangeAttrs , scalar , zeros } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const rangeConfig = { kernelName : Range , backendName : 'tensorflow' , kernelFunc : ( args ) => { const backend = args . backend as NodeJSKernelBackend ; const { start , stop , dtype } = args . attrs as { } as RangeAttrs ; let { step } = args . attrs as { } as RangeAttrs ; const sameStartStop = start === stop ; const increasingRangeNegativeStep = start < stop && step < 0 ; const decreasingRangePositiveStep = stop < start && step > 1 ; if ( sameStartStop || increasingRangeNegativeStep || decreasingRangePositiveStep ) { return zeros ( [ 0 ] , dtype ) ; } if ( stop < start && step === 1 ) { step = - 1 ; } const opAttrs = [ createTensorsTypeOpAttr ( 'Tidx' , dtype ) ] ; const startTensor = scalar ( start ) ; const stopTensor = scalar ( stop ) ; const stepTensor = scalar ( step ) ; const res = backend . executeSingleOutput ( Range , opAttrs , [ startTensor , stopTensor , stepTensor ] ) ; const castedRes = res . cast ( dtype ) ; startTensor . dispose ( ) ; stopTensor . dispose ( ) ; stepTensor . dispose ( ) ; res . dispose ( ) ; return castedRes ; } } ;