license import { ComplexAbs , ComplexAbsInputs , KernelConfig , KernelFunc , TensorInfo } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { ComplexAbsProgram } from '../complex_abs_gpu' ; function makeComplexComponentTensorInfo ( complexTensor , complexPart ) { return { dataId : complexPart . dataId , dtype : complexPart . dtype , shape : complexTensor . shape } ; } export function complexAbs ( args ) { const { inputs , backend } = args ; const { x } = inputs ; const xData = backend . texData . get ( x . dataId ) ; const program = new ComplexAbsProgram ( x . shape ) ; const programInputs = [ makeComplexComponentTensorInfo ( x , xData . complexTensorInfos . real ) , makeComplexComponentTensorInfo ( x , xData . complexTensorInfos . imag ) , ] ; return backend . runWebGLProgram ( program , programInputs , programInputs [ 0 ] . dtype ) ; } export const complexAbsConfig = { kernelName : ComplexAbs , backendName : 'webgl' , kernelFunc : complexAbs as { } as KernelFunc } ;