@genai-fi/nanogpt 0.2.8 → 0.2.10

package/dist/layers/BaseLayer.d.ts

@@ -0,0 +1,8 @@
+import { default as MemoryProfiler } from '../utilities/profile';
+export default abstract class BaseLayer {
+    protected _profiler?: MemoryProfiler;
+    getProfiler(): MemoryProfiler | undefined;
+    setProfiler(value: MemoryProfiler | undefined): void;
+    startMemory(): void;
+    endMemory(label: string): void;
+}

package/dist/layers/BaseLayer.js

@@ -0,0 +1,18 @@
+class t {
+  _profiler;
+  getProfiler() {
+    return this._profiler;
+  }
+  setProfiler(r) {
+    this._profiler = r;
+  }
+  startMemory() {
+    this._profiler?.startMemory();
+  }
+  endMemory(r) {
+    this._profiler?.endMemory(r);
+  }
+}
+export {
+  t as default
+};

package/dist/layers/CausalSelfAttention.d.ts

@@ -1,13 +1,14 @@
 import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from '../config';
 import { default as RoPECache } from './RoPECache';
+import { default as BaseLayer } from './BaseLayer';
 export type KVCache = {
     k: TF.Tensor;
     v: TF.Tensor;
     length: number;
     cumulativeLength: number;
 };
-export default class CausalSelfAttention {
+export default class CausalSelfAttention extends BaseLayer {
     private readonly ropeCache?;
     private config;
     private cAttn;

package/dist/layers/CausalSelfAttention.js

@@ -1,7 +1,8 @@
 import { attentionMask as z } from "../ops/attentionMask.js";
-class j {
+import S from "./BaseLayer.js";
+class C extends S {
   constructor(t, i, s, e) {
-    this.ropeCache = e, this.config = s, this.tf = t, this.index = i, this.cAttn = this.tf.layers.dense({
+    super(), this.ropeCache = e, this.config = s, this.tf = t, this.index = i, this.cAttn = this.tf.layers.dense({
       units: 3 * s.nEmbed,
       useBias: s.biasInLinear,
       name: `block_${i}_attn_cAttn`,
@@ -94,23 +95,24 @@ class j {
     if (e && !this.config.useRope)
       throw new Error("Cannot use pastKV without RoPE enabled");
     return this.tf.tidy(() => {
+      this.startMemory();
       const [o, c, r] = this.getQKV(t), h = o.shape[2], a = this.config.blockSize, u = e ? e.cumulativeLength : 0, [f, d] = this.ropeCache ? this.ropeCache.applyRoPE(o, c, u) : [o, c];
       let n = d, l = r, p = 0;
       e && (p = e.length, n = this.tf.concat([e.k, d], 2), l = this.tf.concat([e.v, r], 2));
       const b = n.shape[2];
       if (b > a) {
-        const k = b - a, g = n.shape[0], I = n.shape[1], _ = n.shape[3];
-        n = n.slice([0, 0, k, 0], [g, I, a, _]), l = l.slice([0, 0, k, 0], [g, I, a, _]), p = a - h;
+        const k = b - a, g = n.shape[0], A = n.shape[1], I = n.shape[3];
+        n = n.slice([0, 0, k, 0], [g, A, a, I]), l = l.slice([0, 0, k, 0], [g, A, a, I]), p = a - h;
       }
       let m;
       p > 0 ? m = this.getAttentionScoresWithPast(f, n, i, p) : m = this.getAttentionScores(f, n, i);
-      const v = this.tf.matMul(m, l), A = this.getOutputProjection(v, i), P = {
+      const _ = this.tf.matMul(m, l), v = this.getOutputProjection(_, i), y = {
         k: this.tf.keep(n),
         v: this.tf.keep(l),
         length: p + h,
         cumulativeLength: e ? e.cumulativeLength + h : h
-      };
-      return { output: A, attention: s ? m.mean(1) : void 0, presentKV: P };
+      }, P = s ? m.mean(1) : void 0;
+      return this.endMemory("CausalSelfAttention"), { output: v, attention: P, presentKV: y };
     });
   }
   dispose() {
@@ -118,5 +120,5 @@ class j {
   }
 }
 export {
-  j as default
+  C as default
 };

package/dist/layers/MLP.d.ts

@@ -1,6 +1,7 @@
 import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from '../config';
-export default class MLP {
+import { default as BaseLayer } from './BaseLayer';
+export default class MLP extends BaseLayer {
     private cFc;
     private cProj;
     private dropout;

package/dist/layers/MLP.js

@@ -1,31 +1,32 @@
-class l {
+import a from "./BaseLayer.js";
+class l extends a {
   cFc;
   cProj;
   dropout;
   tf;
   index;
   _trainable = !0;
-  constructor(t, e, i) {
-    this.tf = t, this.index = e, this.cFc = this.tf.layers.dense({
-      units: i.mlpFactor * i.nEmbed,
+  constructor(t, i, e) {
+    super(), this.tf = t, this.index = i, this.cFc = this.tf.layers.dense({
+      units: e.mlpFactor * e.nEmbed,
       activation: "gelu",
-      useBias: i.biasInLinear,
+      useBias: e.biasInLinear,
       kernelInitializer: this.tf.initializers.randomNormal({
         mean: 0,
         stddev: 0.02
       }),
       biasInitializer: "zeros",
-      name: `block_${e}_mlp_cFc`
+      name: `block_${i}_mlp_cFc`
     }), this.cProj = this.tf.layers.dense({
-      units: i.nEmbed,
-      useBias: i.biasInLinear,
+      units: e.nEmbed,
+      useBias: e.biasInLinear,
       kernelInitializer: this.tf.initializers.randomNormal({
         mean: 0,
-        stddev: 0.02 / Math.sqrt(2 * i.nLayer)
+        stddev: 0.02 / Math.sqrt(2 * e.nLayer)
       }),
       biasInitializer: "zeros",
-      name: `block_${e}_mlp_cProj`
-    }), this.dropout = this.tf.layers.dropout({ rate: i.dropout });
+      name: `block_${i}_mlp_cProj`
+    }), this.dropout = this.tf.layers.dropout({ rate: e.dropout });
   }
   get variables() {
     return [
@@ -45,10 +46,11 @@ class l {
   loadWeights(t) {
     this.cFc.setWeights(t.get(`block_${this.index}_mlpHidden`) || []), this.cProj.setWeights(t.get(`block_${this.index}_mlpOut`) || []);
   }
-  call(t, e = !1) {
+  call(t, i = !1) {
     return this.tf.tidy(() => {
-      const i = this.cFc.apply(t), s = this.cProj.apply(i);
-      return this.dropout.apply(s, { training: e });
+      this.startMemory();
+      const e = this.cFc.apply(t), s = this.cProj.apply(e), r = this.dropout.apply(s, { training: i });
+      return this.endMemory("MLP"), r;
     });
   }
   dispose() {

package/dist/layers/RMSNorm.d.ts

@@ -1,5 +1,6 @@
 import { default as TF } from '@tensorflow/tfjs';
-export default class RMSNorm {
+import { default as BaseLayer } from './BaseLayer';
+export default class RMSNorm extends BaseLayer {
     private gamma;
     private epsilon;
     private tf;

package/dist/layers/RMSNorm.js

@@ -1,26 +1,28 @@
-class m {
+import m from "./BaseLayer.js";
+class o extends m {
   gamma;
   epsilon;
   tf;
-  constructor(a, s, t = 1e-8, e = "") {
-    this.tf = a, this.epsilon = t, this.gamma = a.variable(a.ones(s), !0, `${e}_gamma`, "float32");
+  constructor(t, s, a = 1e-8, e = "") {
+    super(), this.tf = t, this.epsilon = a, this.gamma = t.variable(t.ones(s), !0, `${e}_gamma`, "float32");
   }
   get trainableWeights() {
     return [this.gamma];
   }
-  set trainable(a) {
-    this.gamma.trainable = a;
+  set trainable(t) {
+    this.gamma.trainable = t;
   }
   getWeights() {
     return [this.gamma];
   }
-  setWeights(a) {
-    this.gamma.assign(a[0]);
+  setWeights(t) {
+    this.gamma.assign(t[0]);
   }
-  apply(a) {
+  apply(t) {
     return this.tf.tidy(() => {
-      const t = a.square().mean(-1, !0).add(this.epsilon).rsqrt();
-      return a.mul(t).mul(this.gamma);
+      this.startMemory();
+      const a = t.square().mean(-1, !0).add(this.epsilon).rsqrt(), r = t.mul(a).mul(this.gamma);
+      return this.endMemory("RMSNorm"), r;
     });
   }
   dispose() {
@@ -28,5 +30,5 @@ class m {
   }
 }
 export {
-  m as default
+  o as default
 };

package/dist/layers/TiedEmbedding.js

@@ -1,7 +1,7 @@
-import { o as h, c as i, E as o, y as V, D as X, I as Y, F as Z, N as ee, H as te, J as se, K as ne, O as re, Q as ue, g as L, x as ae, T as A, m as ie, U as oe, u as le, b as q, l as C, V as P, w as U, _ as H } from "../index-DQfEAU9u.js";
-import { s as ce, r as f } from "../sum-B-O33dgG.js";
-import { m } from "../mat_mul-CuHB58-H.js";
-import { c as pe } from "../complex-CeoYJn2o.js";
+import { o as h, d as i, E as o, F as V, H as X, I as Y, J as Z, N as ee, K as te, O as se, Q as ne, T as re, U as ue, i as L, z as ae, V as A, a as ie, W as oe, w as le, f as q, p as C, X as P, y as U, _ as H } from "../index-CWQLouWz.js";
+import { s as ce, r as f } from "../sum-CnIf1YOh.js";
+import { m } from "../mat_mul-4v7St11W.js";
+import { c as pe } from "../complex-x7w5HPOS.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/layers/TransformerBlock.d.ts

@@ -2,7 +2,9 @@ import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from '../config';
 import { KVCache } from './CausalSelfAttention';
 import { default as RoPECache } from './RoPECache';
-export default class Block {
+import { default as MemoryProfiler } from '../utilities/profile';
+import { default as BaseLayer } from './BaseLayer';
+export default class Block extends BaseLayer {
     private ln1;
     private attn;
     private ln2;
@@ -12,6 +14,7 @@ export default class Block {
     private _trainable;
     skipped: boolean;
     constructor(tf: typeof TF, index: number, config: GPTConfig, ropeCache?: RoPECache);
+    setProfiler(value: MemoryProfiler | undefined): void;
     get variables(): TF.Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);

package/dist/layers/TransformerBlock.js

@@ -1,7 +1,8 @@
-import r from "./CausalSelfAttention.js";
+import a from "./CausalSelfAttention.js";
 import o from "./MLP.js";
-import a from "./RMSNorm.js";
-class u {
+import r from "./RMSNorm.js";
+import p from "./BaseLayer.js";
+class f extends p {
   ln1;
   attn;
   ln2;
@@ -11,7 +12,10 @@ class u {
   _trainable = !0;
   skipped = !1;
   constructor(t, i, s, e) {
-    this.tf = t, this.index = i, this.ln1 = new a(t, [s.nEmbed], 1e-8, `block_${this.index}_rms1`), this.attn = new r(this.tf, this.index, s, e), this.ln2 = new a(t, [s.nEmbed], 1e-8, `block_${this.index}_rms2`), this.mlp = new o(this.tf, this.index, s);
+    super(), this.tf = t, this.index = i, this.ln1 = new r(t, [s.nEmbed], 1e-8, `block_${this.index}_rms1`), this.attn = new a(this.tf, this.index, s, e), this.ln2 = new r(t, [s.nEmbed], 1e-8, `block_${this.index}_rms2`), this.mlp = new o(this.tf, this.index, s);
+  }
+  setProfiler(t) {
+    this._profiler = t, this.attn.setProfiler(t), this.mlp.setProfiler(t), this.ln1.setProfiler(t), this.ln2.setProfiler(t);
   }
   get variables() {
     return [
@@ -54,5 +58,5 @@ class u {
   }
 }
 export {
-  u as default
+  f as default
 };

package/dist/{mat_mul-CuHB58-H.js → mat_mul-4v7St11W.js}

@@ -1,4 +1,4 @@
-import { o as c, c as s, b as m, E as M, B as p } from "./index-DQfEAU9u.js";
+import { o as m, d as s, f as c, E as M, B as f } from "./index-CWQLouWz.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -15,13 +15,13 @@ import { o as c, c as s, b as m, E as M, B as p } from "./index-DQfEAU9u.js";
  * limitations under the License.
  * =============================================================================
  */
-function b(e, o, n = !1, l = !1) {
+function p(e, o, n = !1, l = !1) {
   let a = s(e, "a", "matMul"), t = s(o, "b", "matMul");
-  [a, t] = m(a, t);
+  [a, t] = c(a, t);
   const r = { a, b: t }, u = { transposeA: n, transposeB: l };
-  return M.runKernel(p, r, u);
+  return M.runKernel(f, r, u);
 }
-const i = /* @__PURE__ */ c({ matMul_: b });
+const i = /* @__PURE__ */ m({ matMul_: p });
 export {
   i as m
 };

package/dist/ops/attentionMask.js

@@ -1,14 +1,14 @@
-import { engine as d } from "@tensorflow/tfjs";
-import { r as k, s as u } from "../index-DQfEAU9u.js";
-import { m as l } from "../mat_mul-CuHB58-H.js";
-class p {
+import { engine as k } from "@tensorflow/tfjs";
+import { r as m, c as d, s as p } from "../index-CWQLouWz.js";
+import { m as f } from "../mat_mul-4v7St11W.js";
+class h {
   variableNames = ["q", "k", "mask"];
   outputShape;
   userCode;
   // enableShapeUniforms = true;
   customUniforms = [{ name: "divisor", type: "float" }];
-  constructor(t, e, n, a) {
-    this.outputShape = [t, e, n, n], this.userCode = `
+  constructor(e, n, s, a) {
+    this.outputShape = [e, n, s, s], this.userCode = `
         void main() {
             ivec4 coords = getOutputCoords(); // [batch, nh, t1, t2]
             int b = coords.x;
@@ -34,29 +34,55 @@ class p {
         `;
   }
 }
-function f(s) {
-  const { q: t, k: e, mask: n } = s.inputs, { divisor: a } = s.attrs, o = s.backend, c = t.shape[0], i = t.shape[2], r = t.shape[1], m = new p(c, r, i, t.shape[3]);
-  return o.runWebGLProgram(m, [t, e, n], "float32", [[a]]);
+function v(t) {
+  const { q: e, k: n, mask: s } = t.inputs, { divisor: a } = t.attrs, o = t.backend, r = e.shape[0], i = e.shape[2], c = e.shape[1], u = new h(r, c, i, e.shape[3]);
+  return o.runWebGLProgram(u, [e, n, s], "float32", [[a]]);
 }
-const h = {
+const b = {
   kernelName: "AttentionMask",
   backendName: "webgl",
-  kernelFunc: f
+  kernelFunc: v
 };
-k(h);
-function b(s) {
-  const { q: t, k: e, mask: n } = s.inputs, { divisor: a } = s.attrs, o = t.shape[2], i = l(t, e, !1, !0).mul(u(a)), r = n.slice([0, 0], [o, o]).expandDims(0).expandDims(0);
-  return i.add(r);
+m(b);
+function l(t) {
+  const { q: e, k: n, mask: s } = t.inputs, { divisor: a } = t.attrs, o = e.shape[2], i = f(e, n, !1, !0).mul(p(a)), c = s.slice([0, 0], [o, o]).expandDims(0).expandDims(0);
+  return i.add(c);
 }
-const v = {
+const M = {
   kernelName: "AttentionMask",
   backendName: "cpu",
-  kernelFunc: b
+  kernelFunc: l
 };
-k(v);
-function C(s, t, e, n) {
-  return d().runKernel("AttentionMask", { q: s, k: t, mask: e }, { divisor: n });
+m(M);
+const g = {
+  kernelName: "AttentionMask",
+  backendName: "tensorflow",
+  kernelFunc: l
+};
+m(g);
+function N(t, e, n, s) {
+  return k().runKernel("AttentionMask", { q: t, k: e, mask: n }, { divisor: s });
 }
+const A = {
+  kernelName: "AttentionMask",
+  inputsToSave: ["q", "k"],
+  outputsToSave: [],
+  gradFunc: (t, e, n) => {
+    if (Array.isArray(t))
+      throw new Error("Expected dy to be a single Tensor");
+    const [s, a] = e, { divisor: o } = n;
+    return {
+      q: () => t.matMul(a).mul(o),
+      k: () => s.transpose([0, 1, 3, 2]).matMul(t).mul(o).transpose([0, 1, 3, 2]),
+      mask: () => t,
+      divisor: () => {
+        const r = s.matMul(a, !1, !0);
+        return t.mul(r).sum();
+      }
+    };
+  }
+};
+d(A);
 export {
-  C as attentionMask
+  N as attentionMask
 };

package/dist/ops/gatherSub.js

@@ -1,6 +1,6 @@
 import { engine as l } from "@tensorflow/tfjs";
-import { o as g, c as i, E as b, G as d, r as c, a as h } from "../index-DQfEAU9u.js";
-import { r as p, s as f } from "../stack-C9cTkqpq.js";
+import { o as g, d as i, E as b, G as d, r as c, b as h } from "../index-CWQLouWz.js";
+import { r as p, s as f } from "../stack-CTdK-itU.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/ops/node/sparseCrossEntropy.js

@@ -1,4 +1,4 @@
-import { r as o } from "../../index-DQfEAU9u.js";
+import { r as o } from "../../index-CWQLouWz.js";
 function r(e) {
   const { logits: t, labels: n } = e.inputs;
   return e.backend.executeMultipleOutputs("SparseSoftmaxCrossEntropyWithLogits", [], [t, n], 2);

package/dist/ops/scatterSub.js

@@ -1,7 +1,7 @@
 import { engine as $ } from "@tensorflow/tfjs";
-import { j as u, k as S, l as p, E as f, n as E, o as N, c as l, p as y, r as h, a as D, m as x } from "../index-DQfEAU9u.js";
-import { c as m } from "../complex-CeoYJn2o.js";
-import { r as v, s as T } from "../stack-C9cTkqpq.js";
+import { l as u, n as S, p, E as f, q as E, o as N, d as l, t as y, r as h, b as D, a as x } from "../index-CWQLouWz.js";
+import { c as d } from "../complex-x7w5HPOS.js";
+import { r as v, s as T } from "../stack-CTdK-itU.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -21,7 +21,7 @@ import { r as v, s as T } from "../stack-C9cTkqpq.js";
 function i(e, t = "float32") {
   if (u(e), t === "complex64") {
     const a = i(e, "float32"), o = i(e, "float32");
-    return m(a, o);
+    return d(a, o);
   }
   const r = S(p(e), t);
   return f.makeTensor(r, e, t);
@@ -42,10 +42,10 @@ function i(e, t = "float32") {
  * limitations under the License.
  * =============================================================================
  */
-function d(e, t = "float32") {
+function m(e, t = "float32") {
   if (u(e), t === "complex64") {
-    const a = d(e, "float32"), o = i(e, "float32");
-    return m(a, o);
+    const a = m(e, "float32"), o = i(e, "float32");
+    return d(a, o);
   }
   const r = E(p(e), t);
   return f.makeTensor(r, e, t);
@@ -133,7 +133,7 @@ const K = {
 };
 h(K);
 function A(e) {
-  const { logits: t, labels: r, dy: a } = e.inputs, o = r.shape[0], s = t.shape[1], n = v(0, o, 1, "int32"), c = T([n, r], 1), b = d([o]), g = I(c, b, [o, s]), k = D(t, g), w = a.reshape([o, 1]);
+  const { logits: t, labels: r, dy: a } = e.inputs, o = r.shape[0], s = t.shape[1], n = v(0, o, 1, "int32"), c = T([n, r], 1), b = m([o]), g = I(c, b, [o, s]), k = D(t, g), w = a.reshape([o, 1]);
   return x(k, w);
 }
 const F = {
@@ -142,9 +142,9 @@ const F = {
   kernelFunc: A
 };
 h(F);
-function M(e, t, r) {
+function R(e, t, r) {
   return $().runKernel("EfficientScatterSub", { logits: e, labels: t, dy: r }, {});
 }
 export {
-  M as scatterSub
+  R as scatterSub
 };

package/dist/{stack-C9cTkqpq.js → stack-CTdK-itU.js}

@@ -1,4 +1,4 @@
-import { E as e, R as c, o as f, f as u, g as a, P as i } from "./index-DQfEAU9u.js";
+import { E as e, R as c, o as f, h as i, i as a, P as u } from "./index-CWQLouWz.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -38,13 +38,13 @@ function l(n, s, t = 1, r = "float32") {
  * =============================================================================
  */
 function k(n, s = 0) {
-  const t = u(n, "tensors", "stack", "string_or_numeric");
+  const t = i(n, "tensors", "stack", "string_or_numeric");
   a(t.length >= 1, () => "Pass at least one tensor to tf.stack"), t.length > 0 && a(s <= t[0].rank, () => "Axis must be <= rank of the tensor");
   const r = t, o = { axis: s };
-  return e.runKernel(i, r, o);
+  return e.runKernel(u, r, o);
 }
-const h = /* @__PURE__ */ f({ stack_: k });
+const g = /* @__PURE__ */ f({ stack_: k });
 export {
   l as r,
-  h as s
+  g as s
 };

package/dist/{sum-B-O33dgG.js → sum-CnIf1YOh.js}

@@ -1,4 +1,4 @@
-import { o, c as a, E as u, h as i, i as p, S as x } from "./index-DQfEAU9u.js";
+import { o, d as a, E as u, j as p, k as i, S as x } from "./index-CWQLouWz.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -17,7 +17,7 @@ import { o, c as a, E as u, h as i, i as p, S as x } from "./index-DQfEAU9u.js";
  */
 function l(n, t) {
   const s = { x: a(n, "x", "reshape", "string_or_numeric") }, r = { shape: t };
-  return u.runKernel(i, s, r);
+  return u.runKernel(p, s, r);
 }
 const h = /* @__PURE__ */ o({ reshape_: l });
 /**
@@ -38,7 +38,7 @@ const h = /* @__PURE__ */ o({ reshape_: l });
  */
 function m(n, t = null, e = !1) {
   let s = a(n, "x", "sum");
-  s.dtype === "bool" && (s = p(s, "int32"));
+  s.dtype === "bool" && (s = i(s, "int32"));
   const r = { x: s }, c = { axis: t, keepDims: e };
   return u.runKernel(x, r, c);
 }

package/dist/training/AdamExt.js

@@ -1,4 +1,4 @@
-import { A as r, m as c, s as h, a as g, e as o } from "../index-DQfEAU9u.js";
+import { A as r, a as c, s as h, b as g, e as o } from "../index-CWQLouWz.js";
 class u extends r {
   constructor(t, e, s, a, i) {
     super(t, e, s, a), this.config = i, this.startLearningRate = t;

package/dist/training/Trainer.js

@@ -1,8 +1,8 @@
 import { DatasetBuilder as d } from "./DatasetBuilder.js";
-import p from "./AdamExt.js";
-class u {
-  constructor(t, e, s, i = 1e-3) {
-    this.tokenizer = s, this.tf = t, this.model = e, this.learningRate = i, this.resetOptimizer(), this.datasetBuilder = new d(this.tf, s, e.config.blockSize);
+import h from "./AdamExt.js";
+class g {
+  constructor(t, s, e, i = 1e-3) {
+    this.tokenizer = e, this.tf = t, this.model = s, this.learningRate = i, this.resetOptimizer(), this.datasetBuilder = new d(this.tf, e, s.config.blockSize);
   }
   model;
   optimizer;
@@ -25,7 +25,7 @@ class u {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const e = new p(
+    const s = new h(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -37,58 +37,59 @@ class u {
         weightDecay: 0
       }
     );
-    this.optimizer = e;
+    this.optimizer = s;
   }
   printGradients(t) {
-    Object.keys(t).forEach((e) => {
-      const s = t[e];
-      console.log(`${e}:`), console.log(`  Shape: ${s.shape}`), console.log(`  Mean: ${this.tf.mean(s).dataSync()[0]}`), console.log(`  Std: ${this.tf.moments(s).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${this.tf.min(s).dataSync()[0]}`), console.log(`  Max: ${this.tf.max(s).dataSync()[0]}`), console.log(`  Norm: ${this.tf.norm(s).dataSync()[0]}`);
+    Object.keys(t).forEach((s) => {
+      const e = t[s];
+      console.log(`${s}:`), console.log(`  Shape: ${e.shape}`), console.log(`  Mean: ${this.tf.mean(e).dataSync()[0]}`), console.log(`  Std: ${this.tf.moments(e).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${this.tf.min(e).dataSync()[0]}`), console.log(`  Max: ${this.tf.max(e).dataSync()[0]}`), console.log(`  Norm: ${this.tf.norm(e).dataSync()[0]}`);
     });
   }
-  trainStep(t, e = !1, s = !1) {
+  trainStep(t, s = !1, e = !1) {
     return this.tf.tidy(() => {
+      this.model.getProfiler()?.startMemory();
       const { xs: i, ys: a } = t, o = () => {
         const { loss: l, logits: c } = this.model.forward(i, a, !0);
         return c.dispose(), l;
       }, { value: n, grads: r } = this.tf.variableGrads(o);
-      return e || (s && (console.log("-------"), this.printGradients(r), console.log("-------")), this.optimizer.applyGradients(r), this.tf.dispose(r)), n;
+      return s ? this.model.getProfiler()?.endMemory("Training") : (e && (console.log("-------"), this.printGradients(r), console.log("-------")), this.optimizer.applyGradients(r), this.model.getProfiler()?.endMemory("Training"), this.tf.dispose(r)), n;
     });
   }
   dummyPass() {
-    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), e = this.tf.zeros([1, this.model.config.blockSize], "int32");
+    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), s = this.tf.zeros([1, this.model.config.blockSize], "int32");
     try {
-      const s = this.trainStep({ xs: t, ys: e }, !0);
-      s.dataSync(), s.dispose();
-    } catch (s) {
-      console.error("Error during dummy pass:", s);
+      const e = this.trainStep({ xs: t, ys: s }, !0);
+      e.dataSync(), e.dispose();
+    } catch (e) {
+      console.error("Error during dummy pass:", e);
     } finally {
-      t.dispose(), e.dispose();
+      t.dispose(), s.dispose();
     }
   }
-  async trainBatch(t, e) {
+  async trainBatch(t, s) {
     try {
-      const s = this.trainStep(e, !1, !1);
-      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, s.array().then((i) => (t.lastLoss = i, t.losses.push(t.lastLoss), s.dispose(), t.lastLoss));
-    } catch (s) {
-      throw console.error(`Error processing batch at step ${t.step}:`, s), this.tf.dispose(), s;
+      const e = this.trainStep(s, !1, !1);
+      return s.xs.dispose(), s.ys.dispose(), t.step++, t.totalSteps++, e.array().then((i) => (t.lastLoss = i, t.losses.push(t.lastLoss), e.dispose(), t.lastLoss));
+    } catch (e) {
+      throw console.error(`Error processing batch at step ${t.step}:`, e), this.tf.dispose(), e;
     }
   }
-  async createTrainValidationSplit(t, e = 32, s = 0.1) {
-    const i = await this.datasetBuilder.createTextDataset(t, e, 0, 1 - s), a = await this.datasetBuilder.createTextDataset(
+  async createTrainValidationSplit(t, s = 32, e = 0.1) {
+    const i = await this.datasetBuilder.createTextDataset(t, s, 0, 1 - e), a = await this.datasetBuilder.createTextDataset(
       t,
-      e,
-      1 - s,
+      s,
+      1 - e,
       1
     );
     return { trainDataset: i, validationDataset: a };
   }
-  async createDataset(t, e = 32) {
-    return await this.datasetBuilder.createTextDataset(t, e);
+  async createDataset(t, s = 32) {
+    return await this.datasetBuilder.createTextDataset(t, s);
   }
   dispose() {
     this.optimizer && this.optimizer.dispose();
   }
 }
 export {
-  u as default
+  g as default
 };