Function complexAbs

tfjs-backend-webgpu/src/kernels/ComplexAbs.ts:35–50

(
    args: {inputs: ComplexAbsInputs, backend: WebGPUBackend})

}

export function complexAbs(
    args: {inputs: ComplexAbsInputs, backend: WebGPUBackend}): TensorInfo {
  const {inputs, backend} = args;
  const {x} = inputs;

  const xData = backend.tensorMap.get(x.dataId);

  const program = new ComplexAbsProgram(x.shape);
  const programInputs = [
    makeComplexComponentTensorInfo(x, xData.complexTensorInfos.real),
    makeComplexComponentTensorInfo(x, xData.complexTensorInfos.imag),
  ];

  return backend.runWebGPUProgram(
      program, programInputs, programInputs[0].dtype);
}

export const complexAbsConfig: KernelConfig = {
  kernelName: ComplexAbs,