Function step

tfjs-backend-webgpu/src/kernels/Step.ts:24–33

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

import {UnaryOpProgram} from '../unary_op_webgpu';

export function step(
    {inputs, attrs, backend}:
        {inputs: UnaryInputs, attrs: StepAttrs, backend: WebGPUBackend}):
    TensorInfo {
  const {x} = inputs;
  const program =
      new UnaryOpProgram(x.shape, UnaryOpType.STEP, 'stepAlpha : f32,');
  const uniformData = [{type: 'float32', data: [attrs.alpha]}];
  return backend.runWebGPUProgram(program, [x], x.dtype, uniformData);
}

export const stepConfig: KernelConfig = {
  kernelName: Step,