import { a as c, e as G } from "../../index-DSGwv2Yx.js";
import { s as x } from "../../sum-D39FeU5h.js";
class y {
  variableNames = ["x", "meanSquare"];
  outputShape;
  userCode;
  constructor(a, e, o, t = !0) {
    t && this.variableNames.push("gamma"), this.outputShape = [a, e, o], this.userCode = `
        void main() {
            ivec3 coords = getOutputCoords();
            float x = getXAtOutCoords();
            float meanSquare = getMeanSquare(coords.x, coords.y, 0);
            ${t ? "float gamma = getGammaAtOutCoords();" : ""}
            float invRms = inversesqrt(meanSquare + 1e-8);
            float normalized = x * invRms;
            float outVal = normalized ${t ? " * gamma" : ""};
            setOutput(outVal);
        }
        `;
  }
}
function S(r) {
  const { x: a, gamma: e } = r.inputs, o = r.backend, t = a.shape[0], m = a.shape[1], u = a.shape[2], s = a.square().mean(-1, !0), n = new y(t, m, u, e !== void 0);
  return o.runWebGLProgram(n, e ? [a, s, e] : [a, s], "float32");
}
const N = {
  kernelName: "RMSNorm",
  backendName: "webgl",
  kernelFunc: S
}, b = {
  kernelName: "RMSNormNoGamma",
  backendName: "webgl",
  kernelFunc: S
};
c(N);
c(b);
class C {
  variableNames = ["x", "meanSquare", "dyGamma", "dyXMean"];
  outputShape;
  userCode;
  constructor(a, e, o) {
    this.outputShape = [a, e, o], this.userCode = `
        void main() {
            ivec3 coords = getOutputCoords();
            float x = getXAtOutCoords();
            float meanSquare = getMeanSquare(coords.x, coords.y, 0) + 1e-8;
            float dyGamma = getDyGammaAtOutCoords();
            float dyXMean = getDyXMean(coords.x, coords.y, 0) / ${o}.0;
            float invRms = inversesqrt(meanSquare);
            float dx = dyGamma * invRms - x * dyXMean * invRms / meanSquare;
            setOutput(dx);
        }
        `;
  }
}
class k {
  variableNames = ["x", "meanSquare", "dy"];
  outputShape;
  userCode;
  constructor(a, e, o) {
    this.outputShape = [a, e, o], this.userCode = `
        void main() {
            ivec3 coords = getOutputCoords();
            float x = getXAtOutCoords();
            float meanSquare = getMeanSquare(coords.x, coords.y, 0) + 1e-8;
            float dy = getDyAtOutCoords();
            float invRms = inversesqrt(meanSquare);
            float dGamma = dy * (x * invRms);
            setOutput(dGamma);
        }
        `;
  }
}
function M(r) {
  const { dy: a, x: e, gamma: o } = r.inputs, t = r.backend, m = e.shape[0], u = e.shape[1], s = e.shape[2], n = o ? a.mul(o) : a, i = n.mul(e), l = i.sum(-1, !0);
  i.dispose();
  const g = e.square(), d = g.mean(-1, !0);
  g.dispose();
  const h = new C(m, u, s), p = t.runWebGLProgram(h, [e, d, n, l], "float32");
  if (o && n.dispose(), l.dispose(), o) {
    const q = new k(m, u, s), f = t.runWebGLProgram(q, [e, d, a], "float32");
    d.dispose();
    const v = x(G().makeTensorFromTensorInfo(f), [0, 1]);
    return t.disposeIntermediateTensorInfo(f), [p, v];
  } else
    return d.dispose(), [p];
}
const R = {
  kernelName: "RMSNormGrad",
  backendName: "webgl",
  kernelFunc: M
};
c(R);
