@genai-fi/nanogpt 0.6.3 → 0.7.1

package/dist/ops/gelu.js

@@ -1,7 +1,7 @@
-import "../index-BAzbokzv.js";
+import "../index-UdZhlibC.js";
 import "./cpu/gelu.js";
 import "./webgl/gelu.js";
-import { d as e, g as i } from "./grads/gelu.js";
+import { d as e, g as i } from "../gelu-BFwVnd1r.js";
 export {
   e as dGelu,
   i as gelu

package/dist/ops/grads/attentionMask.js

@@ -1,4 +1,4 @@
-import { c as m, s as i } from "../../index-BAzbokzv.js";
+import { i as m, b as i } from "../../index-UdZhlibC.js";
 import { matMulMul as u } from "../matMulMul.js";
 const p = {
   kernelName: "AttentionMask",

package/dist/ops/grads/fusedSoftmax.js

@@ -1,6 +1,6 @@
-import { c as f, a as i, b as l } from "../../index-BAzbokzv.js";
+import { i as f, a as i, c as l } from "../../index-UdZhlibC.js";
 import { mulDrop as g } from "../mulDrop.js";
-import { s as T } from "../../sum-BE_Irnim.js";
+import { s as T } from "../../sum-NLYbiDag.js";
 const Y = {
   kernelName: "FusedSoftmax",
   outputsToSave: [!0],

package/dist/ops/grads/gelu.js

@@ -1,26 +1,5 @@
-import { c as t, e as n } from "../../index-BAzbokzv.js";
-import "../cpu/gelu.js";
-import "../webgl/gelu.js";
-const o = {
-  kernelName: "Gelu",
-  inputsToSave: ["x"],
-  outputsToSave: [],
-  gradFunc: (e, r) => {
-    const [u] = r;
-    return {
-      x: () => a(e, u)
-    };
-  }
-};
-t(o);
-function g(e) {
-  return n().runKernel("Gelu", { x: e });
-}
-function a(e, r) {
-  return n().runKernel("GeluGrad", { dy: e, x: r });
-}
+import "../../index-UdZhlibC.js";
+import { a as e } from "../../gelu-BFwVnd1r.js";
 export {
-  a as d,
-  g,
-  o as geluGradConfig
+  e as geluGradConfig
 };

package/dist/ops/grads/matMulGelu.js

@@ -1,17 +1,17 @@
-import { c as a, e as o } from "../../index-BAzbokzv.js";
-function s(e, n, r) {
+import { i as a, e as o } from "../../index-UdZhlibC.js";
+function i(e, n, r) {
   return o().runKernel("MatMulGeluGrad", { dy: e, x: n, kernel: r });
 }
-const d = {
+const s = {
   kernelName: "MatMulGelu",
   inputsToSave: ["x", "kernel"],
   outputsToSave: [],
   gradFunc: (e, n) => {
-    const [r, t] = n, [u, l] = s(e, r, t);
+    const [r, t] = n, [u, l] = i(e, r, t);
     return {
       x: () => u,
       kernel: () => l
     };
   }
 };
-a(d);
+a(s);

package/dist/ops/grads/normRMS.js

@@ -1,20 +1,20 @@
-import { c as t, e as g } from "../../index-BAzbokzv.js";
-function s(r, a, n) {
+import { i as t, e as g } from "../../index-UdZhlibC.js";
+function i(r, a, n) {
   return g().runKernel("RMSNormGrad", { dy: r, x: a, gamma: n });
 }
-const u = {
+const s = {
   kernelName: "RMSNorm",
   inputsToSave: ["x", "gamma"],
   outputsToSave: [],
   gradFunc: (r, a) => {
-    const [n, e] = a, [m, o] = s(r, n, e);
+    const [n, e] = a, [m, o] = i(r, n, e);
     return {
       x: () => m,
       gamma: () => o
     };
   }
 };
-t(u);
+t(s);
 export {
-  u as normRMSGradConfig
+  s as normRMSGradConfig
 };

package/dist/ops/grads/qkv.js

@@ -1,4 +1,4 @@
-import { c as Q } from "../../index-BAzbokzv.js";
+import { i as Q } from "../../index-UdZhlibC.js";
 const V = {
   kernelName: "QKV",
   inputsToSave: ["x", "kernel"],

package/dist/ops/grads/rope.js

@@ -1,6 +1,6 @@
-import { c as a, e as i } from "../../index-BAzbokzv.js";
+import { i, e as a } from "../../index-UdZhlibC.js";
 function p(n, e, s, o) {
-  return i().runKernel("Rope", { x: n, sin: e, cos: s }, { pastLen: o });
+  return a().runKernel("Rope", { x: n, sin: e, cos: s }, { pastLen: o });
 }
 const c = {
   kernelName: "Rope",
@@ -11,4 +11,4 @@ const c = {
     return t.dispose(), { x: () => r };
   }
 };
-a(c);
+i(c);

package/dist/ops/matMulGelu.js

@@ -1,4 +1,4 @@
-import { e as u } from "../index-BAzbokzv.js";
+import { e as u } from "../index-UdZhlibC.js";
 import "./cpu/matMulGelu.js";
 import "./webgl/matMulGelu.js";
 import "./grads/matMulGelu.js";

package/dist/ops/matMulMul.js

@@ -1,4 +1,4 @@
-import { e as u } from "../index-BAzbokzv.js";
+import { e as u } from "../index-UdZhlibC.js";
 import "./cpu/matMulMul.js";
 import "./webgl/matMulMul.js";
 function m(e, r, t, l = !1, n = !1) {

package/dist/ops/mulDrop.js

@@ -1,4 +1,4 @@
-import { e as t } from "../index-BAzbokzv.js";
+import { e as t } from "../index-UdZhlibC.js";
 import "./cpu/mulDropout.js";
 import "./webgl/mulDropout.js";
 function m(r, o, e, n) {

package/dist/ops/normRMS.js

@@ -1,4 +1,4 @@
-import { e as n } from "../index-BAzbokzv.js";
+import { e as n } from "../index-UdZhlibC.js";
 import "./cpu/normRMS.js";
 import "./webgl/normRMS.js";
 import "./grads/normRMS.js";

package/dist/ops/qkv.js

@@ -1,4 +1,4 @@
-import { e as o } from "../index-BAzbokzv.js";
+import { e as o } from "../index-UdZhlibC.js";
 import "./cpu/qkv.js";
 import "./webgl/qkv.js";
 import "./grads/qkv.js";

package/dist/ops/rope.js

@@ -1,8 +1,8 @@
-import { e as p } from "../index-BAzbokzv.js";
-import "../random_width-DWzaOgrn.js";
-import "../register_all_kernels-BfFCQAqs.js";
+import { e as p } from "../index-UdZhlibC.js";
+import "../random_width-DbSpgl4o.js";
+import "../register_all_kernels-DnbAyBXt.js";
 import "../index-Tf7vU29b.js";
-import "../dataset-pgqp-YfL.js";
+import "../dataset-U3PrjwgU.js";
 import "./cpu/rope.js";
 import "./webgl/rope.js";
 import "./grads/rope.js";

package/dist/ops/scatterSub.js

@@ -1,4 +1,4 @@
-import { e as i } from "../index-BAzbokzv.js";
+import { e as i } from "../index-UdZhlibC.js";
 import "./cpu/scatterSub.js";
 import "./webgl/scatterSub.js";
 function c(t, r, e) {

package/dist/ops/webgl/adamAdjust.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgl/adamAdjust.js

@@ -0,0 +1,50 @@
+import { r as n } from "../../Reshape-CEsEp0AI.js";
+import { f } from "../../index-UdZhlibC.js";
+class v {
+  variableNames = ["moments", "value"];
+  outputShape;
+  userCode;
+  customUniforms = [
+    { name: "invBeta1", type: "float" },
+    { name: "invBeta2", type: "float" },
+    { name: "learningRate", type: "float" },
+    { name: "epsilon", type: "float" }
+  ];
+  constructor(a) {
+    this.outputShape = a, 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;
+
+            setOutput(adjustedValue);
+        }
+        `;
+  }
+}
+function h(t) {
+  const { moments: a, value: s } = t.inputs, { beta1: i, beta2: l, learningRate: u, epsilon: p } = t.attrs, e = t.backend, r = n({ inputs: { x: a }, backend: e, attrs: { shape: [-1] } }), o = n({ inputs: { x: s }, backend: e, attrs: { shape: [-1] } }), d = new v(o.shape), m = e.runWebGLProgram(d, [r, o], "float32", [
+    [1 / i],
+    [1 / l],
+    [u],
+    [p]
+  ]);
+  e.disposeIntermediateTensorInfo(r), e.disposeIntermediateTensorInfo(o);
+  const c = n({ inputs: { x: m }, backend: e, attrs: { shape: s.shape } });
+  return e.disposeIntermediateTensorInfo(m), c;
+}
+const g = {
+  kernelName: "AdamAdjust",
+  backendName: "webgl",
+  kernelFunc: h
+};
+f(g);

package/dist/ops/webgl/adamMoments.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgl/adamMoments.js

@@ -0,0 +1,38 @@
+import { f as m } from "../../index-UdZhlibC.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}]`, s = t === 1 ? "coords" : e.slice(0, -1).map((r, c) => `coords[${c}]`).join(", ");
+    this.userCode = `
+        void main() {
+            float m = getMomentsAtOutCoords();
+            ${n} coords = getOutputCoords();
+            float g = getGradient(${s});
+            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, s = o.backend, r = new i(e.shape);
+  return s.runWebGLProgram(r, [e, t], "float32", [[n], [a]]);
+}
+const g = {
+  kernelName: "AdamMoments",
+  backendName: "webgl",
+  kernelFunc: d
+};
+m(g);

package/dist/ops/webgl/appendCache.js

@@ -1,4 +1,4 @@
-import { r as p } from "../../index-BAzbokzv.js";
+import { f as p } from "../../index-UdZhlibC.js";
 class m {
   variableNames = ["cache", "item"];
   outputShape;

package/dist/ops/webgl/attentionMask.js

@@ -1,4 +1,4 @@
-import { r as m } from "../../index-BAzbokzv.js";
+import { f as m } from "../../index-UdZhlibC.js";
 class h {
   variableNames = ["q", "k"];
   outputShape;

package/dist/ops/webgl/fusedSoftmax.js

@@ -1,7 +1,7 @@
-import { m as b, s as I, r as k } from "../../RealDiv-BYViZwhN.js";
-import { r as v } from "../../Reshape-t7Kcikjk.js";
-import { r as w, p as P } from "../../index-BAzbokzv.js";
-import { e as S } from "../../axis_util-Bu4h7XWV.js";
+import { m as b, s as I, r as k } from "../../RealDiv-CVYNbZxu.js";
+import { r as v } from "../../Reshape-CEsEp0AI.js";
+import { f as w, p as P } from "../../index-UdZhlibC.js";
+import { e as S } from "../../axis_util-5DTW2tFV.js";
 class T {
   variableNames = ["logits", "maxLogits"];
   outputShape;

package/dist/ops/webgl/gatherSub.js

@@ -1,17 +1,17 @@
-import { r as l } from "../../index-BAzbokzv.js";
+import { f as l } from "../../index-UdZhlibC.js";
 class u {
   variableNames = ["labels", "logits", "values"];
   outputShape;
   userCode;
   constructor(e) {
     this.outputShape = [e], this.userCode = `
-      void main() {
-        int coords = getOutputCoords();
-        int index = int(getLabelsAtOutCoords());
-        float val = getValuesAtOutCoords();
-        float logit = getLogits(coords, index);
-        setOutput(val - logit);
-      }
+          void main() {
+              int coords = getOutputCoords();
+              int index = int(getLabelsAtOutCoords());
+              float val = getValuesAtOutCoords();
+              float logit = getLogits(coords, index);
+              setOutput(val - logit);
+          }
     `;
   }
 }

package/dist/ops/webgl/gelu.js

@@ -1,5 +1,5 @@
-import { r as a } from "../../index-BAzbokzv.js";
-import { u as s, C as x } from "../../kernel_funcs_utils-CUxJCg0g.js";
+import { f as a } from "../../index-UdZhlibC.js";
+import { u as s, C as x } from "../../kernel_funcs_utils-CXDy3EN7.js";
 const t = 0.7978845608028654, r = 0.044715, c = x + `
     float x3 = x * x * x;
     float inner = x + ${r} * x3;

package/dist/ops/webgl/log.js

@@ -1,6 +1,6 @@
-import { r, a9 as e } from "../../index-BAzbokzv.js";
-import { u as s, l as N } from "../../kernel_funcs_utils-CUxJCg0g.js";
-import { aG as l } from "../../shared-Ca6iDobD.js";
+import { f as e, af as r } from "../../index-UdZhlibC.js";
+import { u as s, l as N } from "../../kernel_funcs_utils-CXDy3EN7.js";
+import { y as l } from "../../shared-DS5waSIY.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -32,8 +32,8 @@ const a = N + `
   packedOpSnippet: t,
   cpuKernelImpl: l
 }), o = {
-  kernelName: e,
+  kernelName: r,
   backendName: "webgl",
   kernelFunc: n
 };
-r(o);
+e(o);

package/dist/ops/webgl/matMulGelu.js

@@ -1,8 +1,8 @@
-import { r as _, t as R, e as C, g as A, h as N, i as H, u as O } from "../../index-BAzbokzv.js";
-import { r as f } from "../../Reshape-t7Kcikjk.js";
-import { M as U } from "../../mulmat_packed_gpu-D4nKF7Je.js";
-import { m as E } from "../../mat_mul-7121rsJk.js";
-const M = 0.7978845608028654, x = 0.044715, q = `
+import { f as _, t as R, e as C, j as A, k as N, l as H, u as O } from "../../index-UdZhlibC.js";
+import { r as f } from "../../Reshape-CEsEp0AI.js";
+import { M as U } from "../../mulmat_packed_gpu-IGPBp6h9.js";
+import { m as E } from "../../mat_mul-dwmZz69e.js";
+const M = 0.7978845608028654, x = 0.044715, j = `
     vec4 x3 = x * x * x;
     vec4 inner = x + ${x} * x3;
     inner = ${M} * inner;
@@ -10,7 +10,7 @@ const M = 0.7978845608028654, x = 0.044715, q = `
     inner = 0.5 * (1.0 + inner);
     vec4 result = x * inner;
     return result;
-`, z = `
+`, q = `
     vec4 a2 = a * a;
     vec4 a3 = a2 * a;
     vec4 u  = ${M} * (a + ${x} * a3);
@@ -34,7 +34,7 @@ function w({
     i === p,
     () => `Error in matMul: inner shapes (${i}) and (${p}) of Tensors with shapes ${e.shape} and ${t.shape} and transposeA=${s} and transposeB=${n} must match.`
   );
-  const v = s ? [d, i, h] : [d, h, i], S = n ? [m, l, p] : [m, p, l], g = f({ inputs: { x: e }, backend: a, attrs: { shape: v } }), D = f({ inputs: { x: t }, backend: a, attrs: { shape: S } }), G = [g, D], y = Math.max(d, m), L = c, B = O(e.dtype, t.dtype), F = new U(
+  const v = s ? [d, i, h] : [d, h, i], S = n ? [m, l, p] : [m, p, l], k = f({ inputs: { x: e }, backend: a, attrs: { shape: v } }), D = f({ inputs: { x: t }, backend: a, attrs: { shape: S } }), G = [k, D], y = Math.max(d, m), L = c, B = O(e.dtype, t.dtype), F = new U(
     v,
     S,
     [y, h, l],
@@ -44,15 +44,15 @@ function w({
     L,
     !!o,
     !1
-  ), k = [g, D];
-  o && k.push(o);
-  const $ = a.runWebGLProgram(F, k, B), I = f({ inputs: { x: $ }, backend: a, attrs: { shape: b } });
+  ), g = [k, D];
+  o && g.push(o);
+  const $ = a.runWebGLProgram(F, g, B), I = f({ inputs: { x: $ }, backend: a, attrs: { shape: b } });
   G.push($);
   for (const P of G)
     a.disposeIntermediateTensorInfo(P);
   return I;
 }
-function W(e) {
+function z(e) {
   const { inputs: t, backend: s } = e, { x: n, kernel: a } = t;
   if (n === void 0 || a === void 0)
     throw new Error("BatchMatMul requires two input tensors.");
@@ -62,15 +62,15 @@ function W(e) {
     transposeA: !1,
     transposeB: !1,
     backend: s,
-    activationSnippet: q
+    activationSnippet: j
   });
 }
-const j = {
+const W = {
   kernelName: "MatMulGelu",
   backendName: "webgl",
-  kernelFunc: W
+  kernelFunc: z
 };
-_(j);
+_(W);
 function J(e) {
   const { dy: t, x: s, kernel: n } = e.inputs, a = e.backend;
   return R(() => {
@@ -81,7 +81,7 @@ function J(e) {
         transposeA: !1,
         transposeB: !1,
         backend: a,
-        activationSnippet: z,
+        activationSnippet: q,
         multiplier: t
       })
     ), o = E(c, n, !1, !0), r = E(s, c, !0, !1);
@@ -97,5 +97,5 @@ _(Q);
 export {
   te as MATMUL_SHARED_DIM_THRESHOLD,
   w as batchMatMulGeluImpl,
-  W as batchMatMulKernel
+  z as batchMatMulKernel
 };

package/dist/ops/webgl/matMulMul.js

@@ -1,4 +1,4 @@
-import { r as u } from "../../index-BAzbokzv.js";
+import { f as u } from "../../index-UdZhlibC.js";
 import { batchMatMulGeluImpl as c } from "./matMulGelu.js";
 const M = `
     return a * b;

package/dist/ops/webgl/mulDropout.js

@@ -1,4 +1,4 @@
-import { r as m } from "../../index-BAzbokzv.js";
+import { f as m } from "../../index-UdZhlibC.js";
 class f {
   variableNames = ["a", "b"];
   outputShape;
@@ -7,8 +7,8 @@ class f {
     { name: "dropoutRate", type: "float" },
     { name: "seed", type: "float" }
   ];
-  constructor(r, t, o) {
-    this.outputShape = [r, t, o, o], this.userCode = `
+  constructor(t, r, o) {
+    this.outputShape = [t, r, o, o], this.userCode = `
         float random(ivec4 coords) {
             float x = float(coords.x * 4096 + coords.y * 256 + coords.z * 16 + coords.w);
             return fract(sin(seed + x) * 43758.5453123);
@@ -27,7 +27,7 @@ class f {
   }
 }
 function b(e) {
-  const { inputs: r, attrs: t } = e, { a: o, b: s } = r, { dropoutRate: a, seed: c } = t, n = e.backend, d = o.shape[0], u = o.shape[2], p = o.shape[1], l = new f(d, p, u);
+  const { inputs: t, attrs: r } = e, { a: o, b: s } = t, { dropoutRate: a, seed: c } = r, n = e.backend, d = o.shape[0], u = o.shape[2], p = o.shape[1], l = new f(d, p, u);
   return n.runWebGLProgram(l, [o, s], "float32", [
     [a ?? 0],
     [c ?? Math.random() * 1e4]

package/dist/ops/webgl/normRMS.js

@@ -1,5 +1,5 @@
-import { r as p, e as G } from "../../index-BAzbokzv.js";
-import { s as x } from "../../sum-BE_Irnim.js";
+import { f as p, e as G } from "../../index-UdZhlibC.js";
+import { s as x } from "../../sum-NLYbiDag.js";
 class y {
   variableNames = ["x", "meanSquare", "gamma"];
   outputShape;

package/dist/ops/webgl/qkv.js

@@ -1,4 +1,4 @@
-import { r as i } from "../../index-BAzbokzv.js";
+import { f as i } from "../../index-UdZhlibC.js";
 class l {
   variableNames = ["x", "kernel"];
   outputShape;

package/dist/ops/webgl/rope.js

@@ -1,4 +1,4 @@
-import { r as u } from "../../index-BAzbokzv.js";
+import { f as u } from "../../index-UdZhlibC.js";
 class l {
   variableNames = ["x", "sin", "cos"];
   outputShape;

package/dist/ops/webgl/scatterSub.js

@@ -1,4 +1,4 @@
-import { r as i } from "../../index-BAzbokzv.js";
+import { f as i } from "../../index-UdZhlibC.js";
 class u {
   variableNames = ["labels", "softmaxProbs", "dy"];
   outputShape;

package/dist/ops/webgpu/adamAdjust.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/adamAdjust.js

@@ -0,0 +1,52 @@
+import { g as p } from "../../webgpu_program-WaoMq-WD.js";
+import { f as d, c as l } from "../../webgpu_util-DhSeP4b6.js";
+import { f } from "../../index-UdZhlibC.js";
+class c {
+  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];
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = d(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    return `
+        ${p("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 / (1.0 + uniforms.epsilon * invSqrt);
+                let adjustedValue = fma(-uniforms.learningRate * m1Hat, invDenom, value);
+
+                setOutputAtIndex(index, adjustedValue);
+            }
+        }
+    `;
+  }
+}
+function h(t) {
+  const { moments: e, value: a } = t.inputs, { beta1: n, beta2: o, learningRate: i, epsilon: s } = t.attrs, r = t.backend, u = new c(a.shape), m = [
+    { type: "float32", data: [1 / n] },
+    { type: "float32", data: [1 / o] },
+    { type: "float32", data: [i] },
+    { type: "float32", data: [s] }
+  ];
+  return r.runWebGPUProgram(u, [e, a], "float32", m);
+}
+const v = {
+  kernelName: "AdamAdjust",
+  backendName: "webgpu",
+  kernelFunc: h
+};
+f(v);

package/dist/ops/webgpu/adamMoments.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/adamMoments.js

@@ -0,0 +1,51 @@
+import { g as m } from "../../webgpu_program-WaoMq-WD.js";
+import { f as u, c as p } from "../../webgpu_util-DhSeP4b6.js";
+import { f as c } from "../../index-UdZhlibC.js";
+class d {
+  variableNames = ["moments", "gradient"];
+  outputShape;
+  shaderKey = "AdamMoments";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  uniforms = "beta1: f32, beta2: f32";
+  outputComponent = 2;
+  variableComponents = [2, 1];
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = u(this.outputShape.slice(0, -1)), this.dispatch = p(
+      this.dispatchLayout,
+      this.outputShape.slice(0, -1),
+      this.workgroupSize,
+      [1, 1, 1]
+    );
+  }
+  getUserCode() {
+    return `
+        ${m("index")} {
+            if (index < uniforms.size) {
+                let m: vec2<f32> = moments[index];
+                let g: f32 = gradient[index];
+
+                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 f(t) {
+  const { moments: e, gradient: a } = t.inputs, { beta1: n, beta2: o } = t.attrs, s = t.backend, i = new d(e.shape), r = [
+    { type: "float32", data: [n] },
+    { type: "float32", data: [o] }
+  ];
+  return s.runWebGPUProgram(i, [e, a], "float32", r);
+}
+const h = {
+  kernelName: "AdamMoments",
+  backendName: "webgpu",
+  kernelFunc: f
+};
+c(h);

package/dist/ops/webgpu/appendCache.d.ts

@@ -0,0 +1 @@
+export {};