@genai-fi/nanogpt 0.4.4 → 0.5.0

package/dist/ops/webgl/matMulGelu.js

@@ -1,7 +1,7 @@
-import { r as G, t as P, e as R, b as I, n as k, O as L, j as F, Q as U } from "../../index--6vO-cOz.js";
-import { r as g } from "../../Reshape-CiAY8ltP.js";
-import { u as H } from "../../gpgpu_math-CUzjlO9A.js";
-import { m as z } from "../../mat_mul-BEHRPMh0.js";
+import { r as C, t as R, e as I, p as G, N as L, k as F, O as U } from "../../index-iNhkcAEQ.js";
+import { r as S } from "../../Reshape-BE5rA4rT.js";
+import { u as H } from "../../gpgpu_math-C0zyxKFi.js";
+import { m as B } from "../../mat_mul-D0SifYfJ.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -19,39 +19,39 @@ import { m as z } from "../../mat_mul-BEHRPMh0.js";
  * =============================================================================
  */
 class W {
-  constructor(e, s, a, n = !1, c = !1, o = !1, r = null, i = !1, u = !1) {
-    this.variableNames = ["matrixA", "matrixB"], this.packedInputs = !0, this.packedOutput = !0, this.outputShape = a, this.enableShapeUniforms = H(this.outputShape.length);
-    const p = n ? e[1] : e[2], l = Math.ceil(p / 2), b = n ? "i * 2, rc.y" : "rc.y, i * 2", M = c ? "rc.z, i * 2" : "i * 2, rc.z", h = n ? ["a.xxyy", "a.zzww"] : ["a.xxzz", "a.yyww"], d = c ? ["b.xzxz", "b.ywyw"] : ["b.xyxy", "b.zwzw"];
-    let m = "", v = "";
-    r && (i ? m = `vec4 activation(vec4 a) {
+  constructor(e, s, n, a = !1, c = !1, o = !1, r = null, u = !1, l = !1) {
+    this.variableNames = ["matrixA", "matrixB"], this.packedInputs = !0, this.packedOutput = !0, this.outputShape = n, this.enableShapeUniforms = H(this.outputShape.length);
+    const h = a ? e[1] : e[2], p = Math.ceil(h / 2), d = a ? "i * 2, rc.y" : "rc.y, i * 2", $ = c ? "rc.z, i * 2" : "i * 2, rc.z", x = a ? ["a.xxyy", "a.zzww"] : ["a.xxzz", "a.yyww"], m = c ? ["b.xzxz", "b.ywyw"] : ["b.xyxy", "b.zwzw"];
+    let i = "", b = "";
+    r && (u ? i = `vec4 activation(vec4 a) {
           vec4 b = getPreluActivationWeightsAtOutCoords();
           ${r}
-        }` : u ? m = `vec4 activation(vec4 a) {
+        }` : l ? i = `vec4 activation(vec4 a) {
           vec4 b = getLeakyreluAlphaAtOutCoords();
           ${r}
-        }` : m = `vec4 activation(vec4 x) {
+        }` : i = `vec4 activation(vec4 x) {
           ${r}
-        }`, v = "result = activation(result);");
-    const $ = o ? "result += getBiasAtOutCoords();" : "";
-    o && this.variableNames.push("bias"), i && this.variableNames.push("preluActivationWeights"), u && this.variableNames.push("leakyreluAlpha");
-    let f = "rc.x", x = "rc.x";
-    e[0] < s[0] ? f = `imod(rc.x, ${e[0]})` : s[0] < e[0] && (x = `imod(rc.x, ${s[0]})`), this.userCode = `
-      ${m}
+        }`, b = "result = activation(result);");
+    const M = o ? "result += getBiasAtOutCoords();" : "";
+    o && this.variableNames.push("bias"), u && this.variableNames.push("preluActivationWeights"), l && this.variableNames.push("leakyreluAlpha");
+    let f = "rc.x", v = "rc.x";
+    e[0] < s[0] ? f = `imod(rc.x, ${e[0]})` : s[0] < e[0] && (v = `imod(rc.x, ${s[0]})`), this.userCode = `
+      ${i}
       // Don't use uniform for sharedDimensionPacked for performance.
-      const float sharedDimension = ${l}.0;
+      const float sharedDimension = ${p}.0;
 
       vec4 dot2x2ARowBCol(ivec3 rc) {
         vec4 result = vec4(0);
         int batchA = ${f};
-        int batchB = ${x};
-        for (int i = 0; i < ${l}; i++) {
-          vec4 a = getMatrixA(batchA, ${b});
-          vec4 b = getMatrixB(batchB, ${M});
+        int batchB = ${v};
+        for (int i = 0; i < ${p}; i++) {
+          vec4 a = getMatrixA(batchA, ${d});
+          vec4 b = getMatrixB(batchB, ${$});
 
           // These swizzled products need to be separately added.
           // See: https://github.com/tensorflow/tfjs/issues/1735
-          result += (${h[0]} * ${d[0]});
-          result += (${h[1]} * ${d[1]});
+          result += (${x[0]} * ${m[0]});
+          result += (${x[1]} * ${m[1]});
         }
         return result;
       }
@@ -60,97 +60,99 @@ class W {
         ivec3 rc = getOutputCoords();
         vec4 result = dot2x2ARowBCol(rc);
 
-        ${$}
+        ${M}
 
-        ${v}
+        ${b}
 
         setOutput(result);
       }
     `;
   }
 }
-const S = 0.7978845608028654, w = 0.044715, j = `
+const g = 0.7978845608028654, w = 0.044715, j = `
     vec4 x3 = x * x * x;
     vec4 inner = x + ${w} * x3;
-    inner = ${S} * inner;
+    inner = ${g} * inner;
     inner = tanh(inner);
     inner = 0.5 * (1.0 + inner);
     vec4 result = x * inner;
     return result;
 `, q = `
-    vec4 x2 = x * x;
-    vec4 x3 = x2 * x;
-    vec4 u  = ${S} * (x + ${w} * x3);
+    vec4 a2 = a * a;
+    vec4 a3 = a2 * a;
+    vec4 u  = ${g} * (a + ${w} * a3);
     vec4 t  = tanh(u);
     vec4 sech2 = 1.0 - t * t;
-    vec4 du_dx = ${S} * (1.0 + 3.0 * ${w} * x2);
-    vec4 dgelu = 0.5 * (1.0 + t) + 0.5 * x * sech2 * du_dx;
-    return dgelu;
+    vec4 du_dx = ${g} * (1.0 + 3.0 * ${w} * a2);
+    vec4 dgelu = 0.5 * (1.0 + t) + 0.5 * a * sech2 * du_dx;
+    return dgelu * b;
 `, se = 1e3;
-function B({
+function O({
   a: t,
   b: e,
   transposeA: s,
-  transposeB: a,
-  backend: n,
-  activationSnippet: c
+  transposeB: n,
+  backend: a,
+  activationSnippet: c,
+  multiplier: o
 }) {
-  const o = t.shape.length, r = e.shape.length, i = s ? t.shape[o - 2] : t.shape[o - 1], u = a ? e.shape[r - 1] : e.shape[r - 2], p = s ? t.shape[o - 1] : t.shape[o - 2], l = a ? e.shape[r - 2] : e.shape[r - 1], b = t.shape.slice(0, -2), M = e.shape.slice(0, -2), h = k(b), d = k(M), v = L(t.shape.slice(0, -2), e.shape.slice(0, -2)).concat([p, l]);
+  const r = t.shape.length, u = e.shape.length, l = s ? t.shape[r - 2] : t.shape[r - 1], h = n ? e.shape[u - 1] : e.shape[u - 2], p = s ? t.shape[r - 1] : t.shape[r - 2], d = n ? e.shape[u - 2] : e.shape[u - 1], $ = t.shape.slice(0, -2), x = e.shape.slice(0, -2), m = G($), i = G(x), M = L(t.shape.slice(0, -2), e.shape.slice(0, -2)).concat([p, d]);
   F(
-    i === u,
-    () => `Error in matMul: inner shapes (${i}) and (${u}) of Tensors with shapes ${t.shape} and ${e.shape} and transposeA=${s} and transposeB=${a} must match.`
+    l === h,
+    () => `Error in matMul: inner shapes (${l}) and (${h}) of Tensors with shapes ${t.shape} and ${e.shape} and transposeA=${s} and transposeB=${n} must match.`
   );
-  const $ = s ? [h, i, p] : [h, p, i], f = a ? [d, l, u] : [d, u, l], x = g({ inputs: { x: t }, backend: n, attrs: { shape: $ } }), A = g({ inputs: { x: e }, backend: n, attrs: { shape: f } }), y = [x, A], C = Math.max(h, d), O = c, E = U(t.dtype, e.dtype), N = new W(
-    $,
+  const f = s ? [m, l, p] : [m, p, l], v = n ? [i, d, h] : [i, h, d], A = S({ inputs: { x: t }, backend: a, attrs: { shape: f } }), y = S({ inputs: { x: e }, backend: a, attrs: { shape: v } }), k = [A, y], N = Math.max(m, i), E = c, T = U(t.dtype, e.dtype), _ = new W(
     f,
-    [C, p, l],
+    v,
+    [N, p, d],
     s,
-    a,
-    !1,
-    O,
+    n,
     !1,
+    E,
+    !!o,
     !1
-  ), T = [x, A], D = n.runWebGLProgram(N, T, E), _ = g({ inputs: { x: D }, backend: n, attrs: { shape: v } });
-  y.push(D);
-  for (const K of y)
-    n.disposeIntermediateTensorInfo(K);
-  return _;
+  ), D = [A, y];
+  o && D.push(o);
+  const z = a.runWebGLProgram(_, D, T), K = S({ inputs: { x: z }, backend: a, attrs: { shape: M } });
+  k.push(z);
+  for (const P of k)
+    a.disposeIntermediateTensorInfo(P);
+  return K;
 }
-function Q(t) {
-  const { inputs: e, backend: s } = t, { x: a, kernel: n } = e;
-  if (a === void 0 || n === void 0)
+function J(t) {
+  const { inputs: e, backend: s } = t, { x: n, kernel: a } = e;
+  if (n === void 0 || a === void 0)
     throw new Error("BatchMatMul requires two input tensors.");
-  return B({
-    a,
-    b: n,
+  return O({
+    a: n,
+    b: a,
     transposeA: !1,
     transposeB: !1,
     backend: s,
     activationSnippet: j
   });
 }
-const J = {
+const Q = {
   kernelName: "MatMulGelu",
   backendName: "webgl",
-  kernelFunc: Q
+  kernelFunc: J
 };
-G(J);
+C(Q);
 function V(t) {
-  const { dy: e, x: s, kernel: a } = t.inputs, n = t.backend;
-  return P(() => {
-    const c = R().makeTensorFromTensorInfo(
-      B({
+  const { dy: e, x: s, kernel: n } = t.inputs, a = t.backend;
+  return R(() => {
+    const c = I().makeTensorFromTensorInfo(
+      O({
         a: s,
-        b: a,
+        b: n,
         transposeA: !1,
         transposeB: !1,
-        backend: n,
-        activationSnippet: q
+        backend: a,
+        activationSnippet: q,
+        multiplier: e
       })
-    ), o = I(e, c);
-    c.dispose();
-    const r = z(o, a, !1, !0), i = z(s, o, !0, !1);
-    return [r, i];
+    ), o = B(c, n, !1, !0), r = B(s, c, !0, !1);
+    return [o, r];
   });
 }
 const X = {
@@ -158,9 +160,9 @@ const X = {
   backendName: "webgl",
   kernelFunc: V
 };
-G(X);
+C(X);
 export {
   se as MATMUL_SHARED_DIM_THRESHOLD,
-  B as batchMatMulGeluImpl,
-  Q as batchMatMulKernel
+  O as batchMatMulGeluImpl,
+  J as batchMatMulKernel
 };

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

@@ -0,0 +1,14 @@
+import { TensorInfo } from '@tensorflow/tfjs-core';
+import { MathBackendWebGL } from '@tensorflow/tfjs-backend-webgl';
+export declare function batchMatMulKernel(args: {
+    inputs: {
+        x: TensorInfo;
+        kernel: TensorInfo;
+        y: TensorInfo;
+    };
+    attrs: {
+        transposeA: boolean;
+        transposeB: boolean;
+    };
+    backend: MathBackendWebGL;
+}): TensorInfo;

package/dist/ops/webgl/matMulMul.js

@@ -0,0 +1,28 @@
+import { r as u } from "../../index-iNhkcAEQ.js";
+import { batchMatMulGeluImpl as c } from "./matMulGelu.js";
+const M = `
+    return a * b;
+`;
+function p(r) {
+  const { inputs: n, backend: o, attrs: a } = r, { x: t, kernel: e, y: l } = n, { transposeA: i, transposeB: s } = a;
+  if (t === void 0 || e === void 0)
+    throw new Error("BatchMatMul requires two input tensors.");
+  return c({
+    a: t,
+    b: e,
+    transposeA: i,
+    transposeB: s,
+    backend: o,
+    activationSnippet: M,
+    multiplier: l
+  });
+}
+const m = {
+  kernelName: "MatMulMul",
+  backendName: "webgl",
+  kernelFunc: p
+};
+u(m);
+export {
+  p as batchMatMulKernel
+};

package/dist/ops/webgl/mulDropout.js

@@ -1,4 +1,4 @@
-import { r as m } from "../../index--6vO-cOz.js";
+import { r as m } from "../../index-iNhkcAEQ.js";
 class f {
   variableNames = ["a", "b"];
   outputShape;

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

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

package/dist/ops/webgl/normRMS.js

@@ -0,0 +1,86 @@
+import { r as p, e as G } from "../../index-iNhkcAEQ.js";
+import { s as x } from "../../sum-B_92TaHD.js";
+class y {
+  variableNames = ["x", "meanSquare", "gamma"];
+  outputShape;
+  userCode;
+  constructor(a, e, o) {
+    this.outputShape = [a, e, o], this.userCode = `
+        void main() {
+            ivec3 coords = getOutputCoords();
+            float x = getXAtOutCoords();
+            float meanSquare = getMeanSquare(coords.x, coords.y, 0);
+            float gamma = getGammaAtOutCoords();
+            float invRms = inversesqrt(meanSquare + 1e-8);
+            float normalized = x * invRms;
+            float outVal = normalized * gamma;
+            setOutput(outVal);
+        }
+        `;
+  }
+}
+function v(t) {
+  const { x: a, gamma: e } = t.inputs, o = t.backend, r = a.shape[0], n = a.shape[1], m = a.shape[2], u = a.square().mean(-1, !0), s = new y(r, n, m);
+  return o.runWebGLProgram(s, [a, u, e], "float32");
+}
+const C = {
+  kernelName: "RMSNorm",
+  backendName: "webgl",
+  kernelFunc: v
+};
+p(C);
+class b {
+  variableNames = ["x", "meanSquare", "dyGamma", "dyXMean"];
+  outputShape;
+  userCode;
+  constructor(a, e, o) {
+    this.outputShape = [a, e, o], this.userCode = `
+        void main() {
+            ivec3 coords = getOutputCoords();
+            float x = getXAtOutCoords();
+            float meanSquare = getMeanSquare(coords.x, coords.y, 0) + 1e-8;
+            float dyGamma = getDyGammaAtOutCoords();
+            float dyXMean = getDyXMean(coords.x, coords.y, 0) / ${o}.0;
+            float invRms = inversesqrt(meanSquare);
+            float dx = dyGamma * invRms - x * dyXMean * invRms / meanSquare;
+            setOutput(dx);
+        }
+        `;
+  }
+}
+class N {
+  variableNames = ["x", "meanSquare", "dy"];
+  outputShape;
+  userCode;
+  constructor(a, e, o) {
+    this.outputShape = [a, e, o], this.userCode = `
+        void main() {
+            ivec3 coords = getOutputCoords();
+            float x = getXAtOutCoords();
+            float meanSquare = getMeanSquare(coords.x, coords.y, 0) + 1e-8;
+            float dy = getDyAtOutCoords();
+            float invRms = inversesqrt(meanSquare);
+            float dGamma = dy * (x * invRms);
+            setOutput(dGamma);
+        }
+        `;
+  }
+}
+function M(t) {
+  const { dy: a, x: e, gamma: o } = t.inputs, r = t.backend, n = e.shape[0], m = e.shape[1], u = e.shape[2], s = a.mul(o), c = s.mul(e), i = c.sum(-1, !0);
+  c.dispose();
+  const l = e.square(), d = l.mean(-1, !0);
+  l.dispose();
+  const f = new b(n, m, u), S = r.runWebGLProgram(f, [e, d, s, i], "float32");
+  s.dispose(), i.dispose();
+  const h = new N(n, m, u), g = r.runWebGLProgram(h, [e, d, a], "float32");
+  d.dispose();
+  const q = x(G().makeTensorFromTensorInfo(g), [0, 1]);
+  return r.disposeIntermediateTensorInfo(g), [S, q];
+}
+const k = {
+  kernelName: "RMSNormGrad",
+  backendName: "webgl",
+  kernelFunc: M
+};
+p(k);

package/dist/ops/webgl/qkv.js

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

package/dist/ops/webgl/rope.js

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

package/dist/ops/webgl/scatterSub.js

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