license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class BatchNormProgram implements GPGPUProgram { variableNames : string [ ] ; outputShape : number [ ] = [ ] ; userCode : string ; constructor ( xShape , meanShape , varianceShape , offsetShape , scaleShape , varianceEpsilon ) { this . variableNames = [ 'x' , 'mean' , 'variance' ] ; backend_util . assertAndGetBroadcastShape ( xShape , meanShape ) ; backend_util . assertAndGetBroadcastShape ( xShape , varianceShape ) ; let offsetSnippet = '0.0' ; if ( offsetShape != null ) { backend_util . assertAndGetBroadcastShape ( xShape , offsetShape ) ; this . variableNames . push ( 'offset' ) ; offsetSnippet = 'getOffsetAtOutCoords()' ; } let scaleSnippet = '1.0' ; if ( scaleShape != null ) { backend_util . assertAndGetBroadcastShape ( xShape , scaleShape ) ; this . variableNames . push ( 'scale' ) ; scaleSnippet = 'getScaleAtOutCoords()' ; } this . outputShape = xShape ; this . userCode = `
      void main() {
        float x = getXAtOutCoords();
        float mean = getMeanAtOutCoords();
        float variance = getVarianceAtOutCoords();
        float offset = ${ offsetSnippet };
        float scale = ${ scaleSnippet };
        float inv = scale * inversesqrt(variance + float(${ varianceEpsilon }));
        setOutput(dot(vec3(x, -mean, offset), vec3(inv, inv, 1)));
      }
    ` ; } }