@genai-fi/nanogpt 0.2.11 → 0.2.12

package/dist/layers/RoPECache.d.ts

@@ -13,6 +13,5 @@ export default class RoPECache {
     ensureRopeCache(needed: number): void;
     getCos(): TF.Tensor | null;
     getSin(): TF.Tensor | null;
-    applyRoPE(q: TF.Tensor, k: TF.Tensor, pastLen: number): [TF.Tensor, TF.Tensor];
     dispose(): void;
 }

package/dist/layers/RoPECache.js

@@ -1,12 +1,12 @@
-class b {
-  constructor(s, r) {
-    this.tf = s, this.config = r;
-    const o = this.config.nEmbed / this.config.nHead;
-    if (this.rotaryDim = o, this.rotaryDim % 2 !== 0)
+class n {
+  constructor(i, e) {
+    this.tf = i, this.config = e;
+    const t = this.config.nEmbed / this.config.nHead;
+    if (this.rotaryDim = t, this.rotaryDim % 2 !== 0)
       throw new Error("rotaryDim must be even");
     this.ropeBase = 1e4;
-    const i = this.tf.range(0, this.rotaryDim, 2, "float32"), t = i.div(this.tf.scalar(this.rotaryDim, "float32")), e = this.tf.pow(this.tf.scalar(this.ropeBase, "float32"), t);
-    this.ropeInvFreq = this.tf.reciprocal(e), t.dispose(), e.dispose(), i.dispose(), this.config.useRope === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeCacheLen = 0) : this.tf.tidy(() => {
+    const s = this.tf.range(0, this.rotaryDim, 2, "float32"), o = s.div(this.tf.scalar(this.rotaryDim, "float32")), r = this.tf.pow(this.tf.scalar(this.ropeBase, "float32"), o);
+    this.ropeInvFreq = this.tf.reciprocal(r), o.dispose(), r.dispose(), s.dispose(), this.config.useRope === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeCacheLen = 0) : this.tf.tidy(() => {
       this.ensureRopeCache(this.config.blockSize * 4);
     });
   }
@@ -18,11 +18,11 @@ class b {
   ropeSin = null;
   // [cacheLen, rotaryDim/2]
   ropeCacheLen = 0;
-  ensureRopeCache(s) {
-    if (s <= this.ropeCacheLen) return;
+  ensureRopeCache(i) {
+    if (i <= this.ropeCacheLen) return;
     this.ropeCos && this.ropeCos.dispose(), this.ropeSin && this.ropeSin.dispose();
-    const o = this.tf.range(0, s, 1, "float32").expandDims(1).mul(this.ropeInvFreq.expandDims(0));
-    this.ropeCos = this.tf.keep(this.tf.cos(o).expandDims(-1)), this.ropeSin = this.tf.keep(this.tf.sin(o).expandDims(-1)), this.ropeCacheLen = s;
+    const e = Math.max(i, this.ropeCacheLen + this.config.blockSize * 4), s = this.tf.range(0, e, 1, "float32").expandDims(1).mul(this.ropeInvFreq.expandDims(0));
+    this.ropeCos = this.tf.keep(this.tf.cos(s).expandDims(-1)), this.ropeSin = this.tf.keep(this.tf.sin(s).expandDims(-1)), this.ropeCacheLen = e;
   }
   getCos() {
     return this.ropeCos;
@@ -30,21 +30,10 @@ class b {
   getSin() {
     return this.ropeSin;
   }
-  applyRoPE(s, r, o) {
-    const i = s.shape[3], t = this.rotaryDim;
-    if (t > i) return [s, r];
-    const e = s.shape[2], R = o + e;
-    this.ensureRopeCache(R);
-    const n = t / 2, c = this.ropeCos.slice([o, 0, 0], [e, n, 1]).reshape([1, 1, e, n]), a = this.ropeSin.slice([o, 0, 0], [e, n, 1]).reshape([1, 1, e, n]), h = s.shape[0], p = s.shape[1], f = this.tf.range(0, t, 2, "int32"), l = this.tf.range(1, t, 2, "int32"), d = (u) => {
-      const m = u.slice([0, 0, 0, 0], [h, p, e, t]), C = t < i ? u.slice([0, 0, 0, t], [h, p, e, i - t]) : null, g = this.tf.gather(m, f, 3), D = this.tf.gather(m, l, 3), x = g.mul(c).sub(D.mul(a)), k = D.mul(c).add(g.mul(a)), S = this.tf.stack([x, k], -1).reshape([h, p, e, t]);
-      return C ? this.tf.concat([S, C], 3) : S;
-    }, v = d(s), y = d(r);
-    return f.dispose(), l.dispose(), [v, y];
-  }
   dispose() {
     this.ropeCos && this.ropeCos.dispose(), this.ropeSin && this.ropeSin.dispose(), this.ropeInvFreq.dispose();
   }
 }
 export {
-  b as default
+  n as default
 };

package/dist/ops/rope.js

@@ -1,6 +1,6 @@
 import { engine as D } from "@tensorflow/tfjs";
-import { o as G, l as F, k as _, n as z, E as K, p as O, d as T, q as U, r as g, c as A } from "../index-YPKosni4.js";
-import { r as $, s as B } from "../stack-BtKpB0Ry.js";
+import { o as G, l as F, k as _, n as U, E as K, p as z, d as I, q as O, r as f, c as A } from "../index-YPKosni4.js";
+import { r as T, s as B } from "../stack-BtKpB0Ry.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -17,17 +17,17 @@ import { r as $, s as B } from "../stack-BtKpB0Ry.js";
  * limitations under the License.
  * =============================================================================
  */
-function W(e, t = 0) {
-  F(e.length >= 1, () => "Pass at least one tensor to concat");
-  const o = _(e, "tensors", "concat", "string_or_numeric");
-  if (o[0].dtype === "complex64" && o.forEach((s) => {
-    if (s.dtype !== "complex64")
+function W(o, t = 0) {
+  F(o.length >= 1, () => "Pass at least one tensor to concat");
+  const e = _(o, "tensors", "concat", "string_or_numeric");
+  if (e[0].dtype === "complex64" && e.forEach((r) => {
+    if (r.dtype !== "complex64")
       throw new Error(`Cannot concatenate complex64 tensors with a tensor
-          with dtype ${s.dtype}. `);
-  }), o.length === 1)
-    return z(o[0]);
-  const n = o, r = { axis: t };
-  return K.runKernel(O, n, r);
+          with dtype ${r.dtype}. `);
+  }), e.length === 1)
+    return U(e[0]);
+  const n = e, s = { axis: t };
+  return K.runKernel(z, n, s);
 }
 const j = /* @__PURE__ */ G({ concat_: W });
 /**
@@ -46,9 +46,9 @@ const j = /* @__PURE__ */ G({ concat_: W });
  * limitations under the License.
  * =============================================================================
  */
-function H(e, t, o = 0, n = 0) {
-  const r = T(e, "x", "gather"), s = T(t, "indices", "gather", "int32"), c = { x: r, indices: s }, a = { axis: o, batchDims: n };
-  return K.runKernel(U, c, a);
+function H(o, t, e = 0, n = 0) {
+  const s = I(o, "x", "gather"), r = I(t, "indices", "gather", "int32"), c = { x: s, indices: r }, a = { axis: e, batchDims: n };
+  return K.runKernel(O, c, a);
 }
 const E = /* @__PURE__ */ G({ gather_: H });
 class J {
@@ -56,8 +56,9 @@ class J {
   outputShape;
   userCode;
   // enableShapeUniforms = true;
-  constructor(t, o, n, r, s) {
-    this.outputShape = [t, o, n, r], this.userCode = `
+  customUniforms = [{ name: "pastLen", type: "int" }];
+  constructor(t, e, n, s) {
+    this.outputShape = [t, e, n, s], this.userCode = `
         void main() {
             ivec4 coords = getOutputCoords(); // [b, h, t, d]
             int b = coords.x;
@@ -65,14 +66,14 @@ class J {
             int t = coords.z;
             int d = coords.w;
 
-            int rotaryDim = ${r};
+            int rotaryDim = ${s};
 
             float outVal = 0.0;
 
             if (d < rotaryDim) {
                 int pairIdx = d / 2;
-                float cos = getCos(t + ${s}, pairIdx, 0);
-                float sin = getSin(t + ${s}, pairIdx, 0);
+                float cos = getCos(t + pastLen, pairIdx, 0);
+                float sin = getSin(t + pastLen, pairIdx, 0);
 
                 if (d % 2 == 0) {
                     // even index
@@ -95,59 +96,59 @@ class J {
         `;
   }
 }
-function M(e) {
-  const { x: t, sin: o, cos: n } = e.inputs, { pastLen: r } = e.attrs, s = e.backend, c = t.shape[0], a = t.shape[1], i = t.shape[2], d = t.shape[3], p = new J(c, a, i, d, r);
-  return s.runWebGLProgram(p, [t, o, n], "float32");
+function M(o) {
+  const { x: t, sin: e, cos: n } = o.inputs, { pastLen: s } = o.attrs, r = o.backend, c = t.shape[0], a = t.shape[1], i = t.shape[2], d = t.shape[3], p = new J(c, a, i, d);
+  return r.runWebGLProgram(p, [t, e, n], "float32", [[s]]);
 }
 const Q = {
   kernelName: "Rope",
   backendName: "webgl",
   kernelFunc: M
 };
-g(Q);
-function V(e, t, o, n, r) {
-  const s = n.shape[3], c = o;
-  if (c > s) return n;
-  const a = n.shape[2], i = c / 2, d = t.slice([r, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), p = e.slice([r, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), u = n.shape[0], l = n.shape[1], m = $(0, c, 2, "int32"), x = $(1, c, 2, "int32"), X = ((b) => {
-    const v = b.slice([0, 0, 0, 0], [u, l, a, c]), k = c < s ? b.slice([0, 0, 0, c], [u, l, a, s - c]) : null, h = E(v, m, 3), f = E(v, x, 3), C = h.mul(d), y = f.mul(p), R = C.sub(y), N = f.mul(d), S = h.mul(p), w = N.add(S);
-    h.dispose(), f.dispose(), d.dispose(), p.dispose(), C.dispose(), y.dispose(), N.dispose(), S.dispose();
-    const P = B([R, w], -1);
+f(Q);
+function V(o, t, e, n, s) {
+  const r = n.shape[3], c = e;
+  if (c > r) return n;
+  const a = n.shape[2], i = c / 2, d = t.slice([s, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), p = o.slice([s, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), u = n.shape[0], l = n.shape[1], g = T(0, c, 2, "int32"), x = T(1, c, 2, "int32"), $ = ((b) => {
+    const v = b.slice([0, 0, 0, 0], [u, l, a, c]), k = c < r ? b.slice([0, 0, 0, c], [u, l, a, r - c]) : null, h = E(v, g, 3), m = E(v, x, 3), C = h.mul(d), y = m.mul(p), R = C.sub(y), N = m.mul(d), S = h.mul(p), w = N.add(S);
+    h.dispose(), m.dispose(), d.dispose(), p.dispose(), C.dispose(), y.dispose(), N.dispose(), S.dispose();
+    const L = B([R, w], -1);
     R.dispose(), w.dispose();
-    const I = P.reshape([u, l, a, c]);
-    return P.dispose(), k ? j([I, k], 3) : I;
+    const P = L.reshape([u, l, a, c]);
+    return L.dispose(), k ? j([P, k], 3) : P;
   })(n);
-  return m.dispose(), x.dispose(), X;
+  return g.dispose(), x.dispose(), $;
 }
-function L(e) {
-  const { x: t, sin: o, cos: n } = e.inputs, { pastLen: r } = e.attrs, s = t.shape[3];
-  return V(o, n, s, t, r);
+function X(o) {
+  const { x: t, sin: e, cos: n } = o.inputs, { pastLen: s } = o.attrs, r = t.shape[3];
+  return V(e, n, r, t, s);
 }
 const Y = {
   kernelName: "Rope",
   backendName: "cpu",
-  kernelFunc: L
+  kernelFunc: X
 };
-g(Y);
+f(Y);
 const Z = {
   kernelName: "Rope",
   backendName: "tensorflow",
-  kernelFunc: L
+  kernelFunc: X
 };
-g(Z);
-function st(e, t, o) {
-  return t.ensureRopeCache(e.shape[1]), D().runKernel("Rope", { x: e, sin: t.getSin(), cos: t.getCos() }, { pastLen: o });
+f(Z);
+function st(o, t, e) {
+  return t.ensureRopeCache(o.shape[1] + e), D().runKernel("Rope", { x: o, sin: t.getSin(), cos: t.getCos() }, { pastLen: e });
 }
 const q = {
   kernelName: "Rope",
   inputsToSave: ["x", "sin", "cos"],
   outputsToSave: [],
-  gradFunc: (e, t) => {
-    const [o, n, r] = t, s = n.neg(), c = o.shape[3], i = V(s, r, c, e, 0);
-    return s.dispose(), { x: () => i };
+  gradFunc: (o, t) => {
+    const [e, n, s] = t, r = n.neg(), c = e.shape[3], i = V(r, s, c, o, 0);
+    return r.dispose(), { x: () => i };
   }
 };
 A(q);
 export {
   st as rope,
-  L as ropeCPU
+  X as ropeCPU
 };

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.2.11",
+    "version": "0.2.12",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",