license import { backend_util } from '@tensorflow/tfjs-core' ; import { GPGPUProgram } from './gpgpu_math' ; export class BatchNormPackedProgram implements GPGPUProgram { variableNames : string [ ] ; outputShape : number [ ] ; userCode : string ; packedInputs = true ; packedOutput = true ; constructor ( xShape , meanShape , varianceShape , offsetShape , scaleShape , varianceEpsilon ) { this . variableNames = [ 'x' , 'mean' , 'variance' ] ; backend_util . assertAndGetBroadcastShape ( xShape , meanShape ) ; backend_util . assertAndGetBroadcastShape ( xShape , varianceShape ) ; let offsetSnippet = 'vec4(0.0)' ; if ( offsetShape != null ) { backend_util . assertAndGetBroadcastShape ( xShape , offsetShape ) ; this . variableNames . push ( 'offset' ) ; offsetSnippet = 'getOffsetAtOutCoords()' ; } let scaleSnippet = 'vec4(1.0)' ; if ( scaleShape != null ) { backend_util . assertAndGetBroadcastShape ( xShape , scaleShape ) ; this . variableNames . push ( 'scale' ) ; scaleSnippet = 'getScaleAtOutCoords()' ; } this . outputShape = xShape ; this . userCode = `
      void main() {
        vec4 offset = ${ offsetSnippet };
        vec4 scale = ${ scaleSnippet };

        vec4 x = getXAtOutCoords();
        vec4 mean = getMeanAtOutCoords();
        vec4 variance = getVarianceAtOutCoords();

        vec4 inv = scale * inversesqrt(variance + vec4(${ varianceEpsilon }));

        setOutput((x - mean) * inv + offset);
      }
    ` ; } }