license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export const COMPLEX_MULTIPLY = { REAL : 'return areal * breal - aimag * bimag;' , IMAG : 'return areal * bimag + aimag * breal;' } ; export class BinaryOpComplexProgram implements GPGPUProgram { variableNames = [ 'AReal' , 'AImag' , 'BReal' , 'BImag' ] ; userCode : string ; outputShape : number [ ] ; constructor ( op , aShape , bShape ) { this . outputShape = backend_util . assertAndGetBroadcastShape ( aShape , bShape ) ; this . userCode = `
      float binaryOpComplex(
          float areal, float aimag, float breal, float bimag) {
        ${ op }
      }

      void main() {
        float areal = getARealAtOutCoords();
        float aimag = getAImagAtOutCoords();
        float breal = getBRealAtOutCoords();
        float bimag = getBImagAtOutCoords();
        setOutput(binaryOpComplex(areal, aimag, breal, bimag));
      }
    ` ; } }