Method constructor

tfjs-backend-webgl/src/lrn_gpu.ts:25–66

(
      xShape: number[], radius: number, bias: number, alpha: number,
      beta: number)

  userCode: string;

  constructor(
      xShape: number[], radius: number, bias: number, alpha: number,
      beta: number) {
    const rad = radius;
    const maxD = xShape[3] - 1;
    this.outputShape = xShape;

    // optimize pow(bias + alpha * sum, -beta)
    // src: https://github.com/tensorflow/tensorflow/..
    // blob/26033a1644a9c4a5fbe3170ab2e864b6a4ccd4ca/..
    // tensorflow/core/kernels/mkl_lrn_op.cc#L320
    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);
      }
    `;
  }
}