import { e as u } from "../../webgpu_program-CbjdYLYk.js";
import { f as p, c } from "../../webgpu_util-DuofJBMo.js";
import { a as d, a6 as f } from "../../index-DSGwv2Yx.js";
class h {
  variableNames = ["moments", "gradient", "scaling"];
  outputShape;
  shaderKey = "AdamMoments";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "beta1: f32, beta2: f32";
  outputComponent = 2;
  variableComponents = [2, 1, 1];
  constructor(t) {
    this.outputShape = t, this.dispatchLayout = p(this.outputShape.slice(0, -1)), this.dispatch = c(
      this.dispatchLayout,
      this.outputShape.slice(0, -1),
      this.workgroupSize,
      [1, 1, 1]
    );
  }
  getUserCode() {
    return `
        ${u("index")} {
            if (index < uniforms.size) {
                let m: vec2<f32> = moments[index];

                // Loss and clip scaling.
                let g: f32 = gradient[index] * scaling[0];

                let newM1 = fma(m.x, uniforms.beta1, g * (1.0 - uniforms.beta1));
                let newM2 = fma(m.y, uniforms.beta2, g * g * (1.0 - uniforms.beta2));

                setOutputAtIndex(index, vec2<f32>(newM1, newM2));
            }
        }
    `;
  }
}
function l(e) {
  const { moments: t, gradient: a, scaling: s } = e.inputs, { beta1: n, beta2: o } = e.attrs, r = e.backend;
  if (a.dtype !== "float32")
    throw new Error(`Gradient must be float32, but got ${a.dtype}`);
  if (f(t.shape, [...a.shape, 2], "Error in AdamMoments: "), n < 0 || n >= 1)
    throw new Error(`Invalid beta1 value: ${n}. Must be in the range [0, 1).`);
  if (o < 0 || o >= 1)
    throw new Error(`Invalid beta2 value: ${o}. Must be in the range [0, 1).`);
  const i = new h(t.shape), m = [
    { type: "float32", data: [n] },
    { type: "float32", data: [o] }
  ];
  return r.runWebGPUProgram(i, [t, a, s], "float32", m);
}
const g = {
  kernelName: "AdamMoments",
  backendName: "webgpu",
  kernelFunc: l
};
d(g);
