import { a as m } from "../../index-DSGwv2Yx.js";
class i {
  variableNames = ["moments", "gradient"];
  outputShape;
  userCode;
  customUniforms = [
    { name: "beta1", type: "float" },
    { name: "beta2", type: "float" }
  ];
  constructor(e) {
    this.outputShape = e;
    const t = e.length, n = t === 1 ? "int" : `ivec${Math.min(t, 4)}`, a = t === 1 ? "coords" : `coords[${t - 1}]`, r = t === 1 ? "coords" : e.slice(0, -1).map((s, c) => `coords[${c}]`).join(", ");
    this.userCode = `
        void main() {
            float m = getMomentsAtOutCoords();
            ${n} coords = getOutputCoords();

            // Add gradient clipping here
            float g = clamp(getGradient(${r}), -1.0, 1.0);
            int which = ${a};

            float beta = which == 0 ? beta1 : beta2;
            float gg = which == 0 ? g : g * g;
    
            float newM = m * beta + gg * (1.0 - beta);
            setOutput(newM);
        }
        `;
  }
}
function d(o) {
  const { moments: e, gradient: t } = o.inputs, { beta1: n, beta2: a } = o.attrs, r = o.backend, s = new i(e.shape);
  return r.runWebGLProgram(s, [e, t], "float32", [[n], [a]]);
}
const g = {
  kernelName: "AdamMoments",
  backendName: "webgl",
  kernelFunc: d
};
m(g);
