Function step

tfjs-backend-webgl/src/kernels/Step.ts:23–35

(
    {inputs, attrs, backend}:
        {inputs: UnaryInputs, attrs: StepAttrs, backend: MathBackendWebGL})

import {CHECK_NAN_SNIPPET, UnaryOpProgram} from '../unaryop_gpu';

export function step(
    {inputs, attrs, backend}:
        {inputs: UnaryInputs, attrs: StepAttrs, backend: MathBackendWebGL}):
    TensorInfo {
  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: KernelConfig = {
  kernelName: Step,