import { r } from "../../Reshape-BZC-ebeR.js";
import { a as v } from "../../index-DSGwv2Yx.js";
class g {
  variableNames = ["moments", "value"];
  outputShape;
  userCode;
  customUniforms = [
    { name: "invBeta1", type: "float" },
    { name: "invBeta2", type: "float" },
    { name: "learningRate", type: "float" },
    { name: "epsilon", type: "float" }
  ];
  constructor(s, t) {
    this.outputShape = s, t && this.customUniforms.push({ name: "weightDecay", type: "float" }), this.userCode = `
        void main() {
            float v = getValueAtOutCoords();
            int coords = getOutputCoords();
            coords *= 2;
            float m1 = getMoments(coords);
            float m2 = getMoments(coords + 1);

            float m1Hat = m1 * invBeta1;
            float m2Hat = m2 * invBeta2;

            float invSqrt = inversesqrt(max(m2Hat, 1e-30));
            float invDenom = invSqrt / (1.0 + epsilon * invSqrt);
            float adjustedValue = -learningRate * m1Hat * invDenom + v;

            ${t ? "adjustedValue = adjustedValue - learningRate * weightDecay * v;" : ""}

            setOutput(adjustedValue);
        }
        `;
  }
}
function b(a) {
  const { moments: s, value: t } = a.inputs, { beta1: l, beta2: d, learningRate: p, epsilon: c, weightDecay: n } = a.attrs, e = a.backend, i = r({ inputs: { x: s }, backend: e, attrs: { shape: [-1] } }), o = r({ inputs: { x: t }, backend: e, attrs: { shape: [-1] } }), f = new g(o.shape, n > 0), m = [[1 / l], [1 / d], [p], [c]];
  n > 0 && m.push([n]);
  const u = e.runWebGLProgram(f, [i, o], "float32", m);
  e.disposeIntermediateTensorInfo(i), e.disposeIntermediateTensorInfo(o);
  const h = r({ inputs: { x: u }, backend: e, attrs: { shape: t.shape } });
  return e.disposeIntermediateTensorInfo(u), h;
}
const j = {
  kernelName: "AdamAdjust",
  backendName: "webgl",
  kernelFunc: b
};
v(j);
