license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram , useShapeUniforms } from './gpgpu_math' ; export const CHECK_NAN_SNIPPET = `
  if (isnan(a)) return a;
  if (isnan(b)) return b;
` ; export const SQUARED_DIFFERENCE = 'return (a - b) * (a - b);' ; export class BinaryOpProgram implements GPGPUProgram { variableNames = [ 'A' , 'B' ] ; outputShape : number [ ] ; userCode : string ; enableShapeUniforms : boolean ; constructor ( op , aShape , bShape ) { this . outputShape = backend_util . assertAndGetBroadcastShape ( aShape , bShape ) ; this . enableShapeUniforms = useShapeUniforms ( this . outputShape . length ) ; this . userCode = `
      float binaryOperation(float a, float b) {
        ${ op }
      }

      void main() {
        float a = getAAtOutCoords();
        float b = getBAtOutCoords();
        setOutput(binaryOperation(a, b));
      }
    ` ; } }