license import { backend_util , fill , greater , isNaN as tfIsNan , KernelConfig , ones , scalar , Step , StepAttrs , StepInputs , Tensor , tidy , where } from '@tensorflow/tfjs' ; import { createTensorsTypeOpAttr , NodeJSKernelBackend } from '../nodejs_kernel_backend' ; export const stepConfig = { kernelName : Step , backendName : 'tensorflow' , kernelFunc : ( args ) => { const { x } = args . inputs as StepInputs ; const backend = args . backend as NodeJSKernelBackend ; const { alpha } = args . attrs as { } as StepAttrs ; const dtype = x . dtype ; return tidy ( ( ) => { const nans = tfIsNan ( x as Tensor ) ; const stepNoNans = where ( greater ( x as Tensor , scalar ( 0 , dtype ) ) , ones ( x . shape ) , fill ( x . shape , alpha , dtype ) ) ; const opAttrs = [ createTensorsTypeOpAttr ( 'T' , backend_util . upcastType ( x . dtype , stepNoNans . dtype ) ) ] ; return backend . executeSingleOutput ( 'Select' , opAttrs , [ nans , x , stepNoNans ] ) ; } ) ; } } ;