license import { KernelConfig , KernelFunc , Step , StepAttrs , TensorInfo , UnaryInputs } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { CHECK_NAN_SNIPPET , UnaryOpProgram } from '../unaryop_gpu' ; export function step ( { inputs , attrs , backend } ) { const { x } = inputs ; const opSnippet = CHECK_NAN_SNIPPET + `
    return x > 0.0 ? 1.0 : float(${ attrs . alpha });
  ` ; const program = new UnaryOpProgram ( x . shape , opSnippet ) ; return backend . runWebGLProgram ( program , [ x ] , x . dtype ) ; } export const stepConfig = { kernelName : Step , backendName : 'webgl' , kernelFunc : step as { } as KernelFunc , } ;