import { e as d } from "../../webgpu_program-CbjdYLYk.js";
import { f as l, c as f } from "../../webgpu_util-DuofJBMo.js";
import { a as h, a6 as c } from "../../index-DSGwv2Yx.js";
class v {
  variableNames = ["moments", "value"];
  outputShape;
  shaderKey = "AdamAdjust";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "invbeta1: f32, invbeta2: f32, learningRate: f32, epsilon: f32";
  outputComponent = 1;
  variableComponents = [2, 1];
  useWeightDecay;
  constructor(a, t) {
    this.outputShape = a, this.dispatchLayout = l(this.outputShape), this.dispatch = f(this.dispatchLayout, this.outputShape, this.workgroupSize), this.useWeightDecay = t, t && (this.uniforms += ", weightDecay: f32");
  }
  getUserCode() {
    return `
        ${d("index")} {
            if (index < uniforms.size) {
                let moments: vec2<f32> = moments[index];
                let value: f32 = value[index];

                let m1Hat = moments.x * uniforms.invbeta1;
                let m2Hat = moments.y * uniforms.invbeta2;

                let invSqrt = inverseSqrt(max(m2Hat, 1e-30));
                let invDenom = invSqrt / fma(uniforms.epsilon, invSqrt, 1.0);
                var adjustedValue = fma(-uniforms.learningRate * m1Hat, invDenom, value);

                ${this.useWeightDecay ? "adjustedValue = adjustedValue - uniforms.learningRate * uniforms.weightDecay * value;" : ""}

                setOutputAtIndex(index, adjustedValue);
            }
        }
    `;
  }
}
function g(e) {
  const { moments: a, value: t } = e.inputs, { beta1: i, beta2: o, learningRate: u, epsilon: r, weightDecay: s } = e.attrs, m = e.backend;
  c(a.shape, [...t.shape, 2], "Error in AdamAdjust: ");
  const p = new v(t.shape, s > 0), n = [
    { type: "float32", data: [1 / i] },
    { type: "float32", data: [1 / o] },
    { type: "float32", data: [u] },
    { type: "float32", data: [r] }
  ];
  return s > 0 && n.push({ type: "float32", data: [s] }), m.runWebGPUProgram(p, [a, t], "float32", n);
}
const y = {
  kernelName: "AdamAdjust",
  backendName: "webgpu",
  kernelFunc: g
};
h(y);
