license import { GPGPUProgram } from './gpgpu_math' ; export class LRNProgram implements GPGPUProgram { variableNames = [ 'x' ] ; outputShape : number [ ] = [ ] ; userCode : string ; constructor ( xShape , radius , bias , alpha , beta ) { const rad = radius ; const maxD = xShape [ 3 ] - 1 ; this . outputShape = xShape ; let powOperator ; const basis = `float(${ bias }) + float(${ alpha }) * sum` ; if ( beta === 0.5 ) { powOperator = `inversesqrt(${ basis })` ; } else if ( beta === 1.0 ) { powOperator = `1.0/(${ basis })` ; } else { powOperator = `exp(log(${ basis }) * float(-${ beta }));` ; } this . userCode = `
      void main() {
        ivec4 coords = getOutputCoords();
        int b = coords[0];
        int r = coords[1];
        int c = coords[2];
        int d = coords[3];
        float x = getX(b, r, c, d);
        float sum = 0.0;
        for (int j = -${ rad }; j <= ${ rad }; j++) {
          int idx = d + j;
          if (idx >= 0 && idx <=  ${ maxD }) {
            float z = getX(b, r, c, idx);
            sum += z * z;
          }
        }
        float val = x * ${ powOperator };
        setOutput(val);
      }
    ` ; } }