@genai-fi/nanogpt 0.6.2 → 0.7.0

package/dist/{sin-D_h-qCSx.js → sin-Cocju-BY.js}

@@ -1,4 +1,4 @@
-import { o, q as t, E as c, a0 as a, a1 as e } from "./index-BAzbokzv.js";
+import { x as o, y as t, E as c, a3 as a, a4 as x } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -15,11 +15,11 @@ import { o, q as t, E as c, a0 as a, a1 as e } from "./index-BAzbokzv.js";
  * limitations under the License.
  * =============================================================================
  */
-function i(n) {
+function e(n) {
   const s = { x: t(n, "x", "cos", "float32") };
   return c.runKernel(a, s);
 }
-const f = /* @__PURE__ */ o({ cos_: i });
+const f = /* @__PURE__ */ o({ cos_: e });
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -36,11 +36,11 @@ const f = /* @__PURE__ */ o({ cos_: i });
  * limitations under the License.
  * =============================================================================
  */
-function x(n) {
+function i(n) {
   const s = { x: t(n, "x", "sin", "float32") };
-  return c.runKernel(e, s);
+  return c.runKernel(x, s);
 }
-const p = /* @__PURE__ */ o({ sin_: x });
+const p = /* @__PURE__ */ o({ sin_: i });
 export {
   f as c,
   p as s

package/dist/{softmax-fsdtf6JC.js → softmax-GPNK3o-U.js}

@@ -1,4 +1,4 @@
-import { o as r, q as f, E as e, S as i } from "./index-BAzbokzv.js";
+import { x as r, y as f, E as e, S as i } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -22,7 +22,7 @@ function l(s, o = -1) {
   const n = { logits: t }, a = { dim: o };
   return e.runKernel(i, n, a);
 }
-const p = /* @__PURE__ */ r({ softmax_: l });
+const c = /* @__PURE__ */ r({ softmax_: l });
 export {
-  p as s
+  c as s
 };

package/dist/{split-eiktj-6L.js → split-CHzJjxDv.js}

@@ -1,4 +1,4 @@
-import { o as p, q as i, E as a, F as c } from "./index-BAzbokzv.js";
+import { x as p, y as i, E as a, J as c } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -15,9 +15,9 @@ import { o as p, q as i, E as a, F as c } from "./index-BAzbokzv.js";
  * limitations under the License.
  * =============================================================================
  */
-function e(t, s, o = 0) {
-  const n = { x: i(t, "x", "split") }, r = { numOrSizeSplits: s, axis: o };
-  return a.runKernel(c, n, r);
+function e(t, s, n = 0) {
+  const o = { x: i(t, "x", "split") }, r = { numOrSizeSplits: s, axis: n };
+  return a.runKernel(c, o, r);
 }
 const u = /* @__PURE__ */ p({ split_: e });
 export {

package/dist/{stack-dfEEz2OY.js → stack-Dpgg_1W1.js}

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

package/dist/{sum-BE_Irnim.js → sum-B8wEpKsg.js}

@@ -1,4 +1,4 @@
-import { o as e, q as u, L as c, E as l, N as m } from "./index-BAzbokzv.js";
+import { x as e, y as u, Q as c, E as l, U as m } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -15,13 +15,13 @@ import { o as e, q as u, L as c, E as l, N as m } from "./index-BAzbokzv.js";
  * limitations under the License.
  * =============================================================================
  */
-function i(t, o = null, n = !1) {
+function i(t, n = null, o = !1) {
   let s = u(t, "x", "sum");
   s.dtype === "bool" && (s = c(s, "int32"));
-  const r = { x: s }, a = { axis: o, keepDims: n };
+  const r = { x: s }, a = { axis: n, keepDims: o };
   return l.runKernel(m, r, a);
 }
-const f = /* @__PURE__ */ e({ sum_: i });
+const x = /* @__PURE__ */ e({ sum_: i });
 export {
-  f as s
+  x as s
 };

package/dist/{tensor-Xyi595sG.js → tensor-RvZVNmg0.js}

@@ -1,4 +1,4 @@
-import { O as t, Q as a } from "./index-BAzbokzv.js";
+import { V as t, W as a } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/{tensor2d-CPEkynbH.js → tensor2d-B_kyod7_.js}

@@ -1,4 +1,4 @@
-import { V as t, O as s, Q as a } from "./index-BAzbokzv.js";
+import { X as t, V as s, W as a } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/training/AdamExt.js

@@ -1,4 +1,4 @@
-import { A as r, a as c, s as h, b as g, e as o } from "../index-BAzbokzv.js";
+import { A as r, a as c, b as h, c as g, e as o } from "../index-C0dhsYom.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/DatasetBuilder.js

@@ -1,5 +1,5 @@
-import { t as u } from "../index-BAzbokzv.js";
-import { d as z, i as f } from "../dataset-pgqp-YfL.js";
+import { t as u } from "../index-C0dhsYom.js";
+import { d as z, i as f } from "../dataset-JDyjG3QR.js";
 import "../index-Tf7vU29b.js";
 /**
  * @license

package/dist/training/Evaluator.js

@@ -11,7 +11,7 @@ class p {
       if (e.done) break;
       const c = e.value, { xs: r, ys: i } = c, [u, l] = this.model.forward({ training: !1 }, r, i);
       u.dispose(), r.dispose(), i.dispose();
-      const d = l.arraySync();
+      const d = await l.array();
       l.dispose(), t += d, o++;
     }
     return t / o;

package/dist/training/FullTrainer.js

@@ -1,7 +1,7 @@
 import { generateText as T } from "../utilities/generate.js";
 import L from "./Trainer.js";
 import x from "./Evaluator.js";
-import { d as h } from "../index-BAzbokzv.js";
+import { d as h } from "../index-C0dhsYom.js";
 import y from "../utilities/profile.js";
 const D = {
   desiredLoss: 0.01,
@@ -9,12 +9,12 @@ const D = {
   maxSteps: 1e3
 };
 class I extends L {
-  constructor(i, e, o = 3e-4) {
-    super(i, e, o);
+  constructor(i, e, r = 3e-4) {
+    super(i, e, r);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
-  async trainOnDataset(i, e, o) {
-    const { desiredLoss: p, logInterval: g, onStep: l, prompt: c, maxSteps: u } = {
+  async trainOnDataset(i, e, r) {
+    const { desiredLoss: p, logInterval: g, onStep: l, prompt: m, maxSteps: u } = {
       ...D,
       ...e
     }, n = Date.now(), t = {
@@ -27,33 +27,33 @@ class I extends L {
       trainingDuration: 0,
       ...this.lastState || {}
     };
-    this.lastState = t, this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || (this.model.config.layerConfig.profiler = new y())), this.running = !0, t.logStartTime = n;
-    const m = o ? new x(this.model, o) : void 0, f = await i.iterator();
+    this.lastState = t, await this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || (this.model.config.layerConfig.profiler = new y())), this.running = !0, t.logStartTime = n;
+    const c = r ? new x(this.model, r) : void 0, f = await i.iterator();
     try {
       for (; this.running && !(t.lastLoss < p); ) {
-        const r = await f.next();
-        if (r.done) break;
-        const d = r.value, v = this.trainBatch(t, d, e.advancedMetrics || !1), s = {
+        const o = await f.next();
+        if (o.done) break;
+        const d = o.value, S = this.trainBatch(t, d), s = {
           loss: t.lastLoss,
           step: t.step,
           time: Date.now() - n,
           batchSize: d.xs.shape[0],
-          learningRate: e?.advancedMetrics ? this.optimizer.lr : void 0,
-          gradientNorm: e?.advancedMetrics ? t.gradientNorm : void 0
+          learningRate: e?.advancedMetrics ? this.optimizer.lr : void 0
+          //gradientNorm: options?.advancedMetrics ? await state.gradientNorm : undefined,
         };
         if (this.model.log.push(s), t.step % g === 0) {
-          await v;
-          const S = Date.now();
-          if (t.trainingDuration += S - t.logStartTime, m)
+          await S;
+          const v = Date.now();
+          if (t.trainingDuration += v - t.logStartTime, c)
             try {
-              const a = await m.evaluate(5);
+              const a = await c.evaluate(5);
               t.validationLosses.push(a), s.valLoss = a;
             } catch (a) {
               console.error("Validation error:", a);
             }
           if (l) {
-            if (c) {
-              const w = await T(this.tokenizer, this.model, c, 100, {
+            if (m) {
+              const w = await T(this.tokenizer, this.model, m, 100, {
                 temperature: 0.8
               });
               s.example = w;
@@ -70,8 +70,8 @@ class I extends L {
         }
         t.step >= u && this.stop();
       }
-    } catch (r) {
-      throw console.error("Training error:", r), h(), r;
+    } catch (o) {
+      throw console.error("Training error:", o), h(), o;
     }
     return h(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }

package/dist/training/Trainer.d.ts

@@ -11,7 +11,7 @@ export interface TrainingState {
     totalSteps: number;
     losses: number[];
     validationLosses: number[];
-    gradientNorm?: number;
+    gradientNorm?: Promise<number>;
 }
 export interface TrainingProgress {
     duration: number;
@@ -47,16 +47,15 @@ export default abstract class GPTTrainer {
     stop(): void;
     getOptimizer(): AdamExt;
     resetOptimizer(config?: AdamConfig): void;
-    private maxGradNorm;
-    protected trainStep(state: Partial<TrainingState>, batch: {
+    protected trainStep(_state: Partial<TrainingState>, batch: {
         xs: Tensor;
         ys: Tensor;
-    }, dummy?: boolean, calcNorm?: boolean): Scalar;
-    protected dummyPass(): void;
+    }, dummy?: boolean): Scalar;
+    protected dummyPass(): Promise<void>;
     protected trainBatch(state: TrainingState, batch: {
         xs: Tensor;
         ys: Tensor;
-    }, calcNorm?: boolean): Promise<number>;
+    }): Promise<number>;
     abstract trainOnDataset(dataset: Dataset<{
         xs: Tensor;
         ys: Tensor;

package/dist/training/Trainer.js

@@ -1,11 +1,10 @@
-import { DatasetBuilder as g, flattenTokens as m, PAGE_FACTOR as u } from "./DatasetBuilder.js";
-import f from "./AdamExt.js";
-import { t as y, v as z, d as c } from "../index-BAzbokzv.js";
-import { n as S } from "../norm-CzltS9Fz.js";
-import { z as p } from "../zeros-YJDE7oRb.js";
-class R {
-  constructor(t, e, s = 1e-3) {
-    this.tokenizer = e, this.model = t, this.learningRate = s, this.resetOptimizer(), this.datasetBuilder = new g(e, t.config.gpt.blockSize);
+import { DatasetBuilder as m, flattenTokens as p, PAGE_FACTOR as u } from "./DatasetBuilder.js";
+import g from "./AdamExt.js";
+import { t as f, v as y, d as c } from "../index-C0dhsYom.js";
+import { z as h } from "../zeros-DCPCdFGq.js";
+class x {
+  constructor(t, s, e = 1e-3) {
+    this.tokenizer = s, this.model = t, this.learningRate = e, this.resetOptimizer(), this.datasetBuilder = new m(s, t.config.gpt.blockSize);
   }
   model;
   optimizer;
@@ -27,7 +26,7 @@ class R {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const e = new f(
+    const s = new g(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -39,78 +38,78 @@ class R {
         weightDecay: 0
       }
     );
-    this.optimizer = e;
+    this.optimizer = s;
   }
-  maxGradNorm(t) {
-    let e = 0;
-    return Object.keys(t).forEach((s) => {
-      const a = t[s], r = S(a), i = r.dataSync()[0];
-      r.dispose(), i > e && (e = i);
-    }), e;
-  }
-  trainStep(t, e, s = !1, a = !1) {
-    return y(() => {
+  /*private async maxGradNorm(grads: NamedVariableMap): Promise<number> {
+      let maxNorm = 0;
+      // Print all gradients
+      await Promise.all(
+          Object.keys(grads).map(async (varName) => {
+              const grad = grads[varName];
+              const temp = norm(grad);
+              const gradNorm = (await temp.data())[0];
+              temp.dispose();
+              if (gradNorm > maxNorm) {
+                  maxNorm = gradNorm;
+              }
+          })
+      );
+      return maxNorm;
+  }*/
+  trainStep(t, s, e = !1) {
+    return f(() => {
       this.model.getProfiler()?.startMemory();
-      const { xs: r, ys: i } = e, d = () => {
-        const [n, h] = this.model.forward({ training: !0 }, r, i);
-        return n.dispose(), h;
-      }, { value: l, grads: o } = z(d);
-      if (s)
-        this.model.getProfiler()?.endMemory("Training");
-      else {
-        if (a) {
-          const n = this.maxGradNorm(o);
-          t.gradientNorm = n;
-        }
-        this.optimizer.applyGradients(o), this.model.getProfiler()?.endMemory("Training"), c(o);
-      }
-      return l;
+      const { xs: a, ys: i } = s, o = () => {
+        const [l, d] = this.model.forward({ training: !0 }, a, i);
+        return l.dispose(), d;
+      }, { value: n, grads: r } = y(o);
+      return e ? this.model.getProfiler()?.endMemory("Training") : (this.optimizer.applyGradients(r), this.model.getProfiler()?.endMemory("Training"), c(r)), n;
     });
   }
-  dummyPass() {
-    const t = p([1, this.model.config.gpt.blockSize], "int32"), e = p([1, this.model.config.gpt.blockSize], "int32");
+  async dummyPass() {
+    const t = h([1, this.model.config.gpt.blockSize], "int32"), s = h([1, this.model.config.gpt.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);
+      await e.data(), e.dispose();
+    } catch (e) {
+      console.error("Error during dummy pass:", e);
     } finally {
-      t.dispose(), e.dispose();
+      t.dispose(), s.dispose();
     }
   }
-  async trainBatch(t, e, s = !1) {
+  async trainBatch(t, s) {
     try {
-      const a = this.trainStep(t, e, !1, s);
-      return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, a.array().then((r) => (t.lastLoss = r, t.losses.push(t.lastLoss), a.dispose(), t.lastLoss));
-    } catch (a) {
-      throw console.error(`Error processing batch at step ${t.step}:`, a), c(), a;
+      const e = this.trainStep(t, s, !1);
+      return s.xs.dispose(), s.ys.dispose(), t.step++, t.totalSteps++, e.array().then((a) => (t.lastLoss = a, t.losses.push(t.lastLoss), e.dispose(), t.lastLoss));
+    } catch (e) {
+      throw console.error(`Error processing batch at step ${t.step}:`, e), c(), e;
     }
   }
-  async createTrainValidationSplit(t, e = 32, s = 0.1) {
-    const a = await m(t, this.tokenizer), r = /* @__PURE__ */ new Set();
-    if (s > 0) {
-      const l = Math.floor(a.length / (this.datasetBuilder.blockSize * u)), o = Math.max(1, Math.floor(l * s));
-      for (; r.size < o; ) {
-        const n = Math.floor(Math.random() * l);
-        r.add(n);
+  async createTrainValidationSplit(t, s = 32, e = 0.1) {
+    const a = await p(t, this.tokenizer), i = /* @__PURE__ */ new Set();
+    if (e > 0) {
+      const r = Math.floor(a.length / (this.datasetBuilder.blockSize * u)), l = Math.max(1, Math.floor(r * e));
+      for (; i.size < l; ) {
+        const d = Math.floor(Math.random() * r);
+        i.add(d);
       }
     }
-    const i = await this.datasetBuilder.createTextDataset(a, e, r, !1), d = await this.datasetBuilder.createTextDataset(
+    const o = await this.datasetBuilder.createTextDataset(a, s, i, !1), n = await this.datasetBuilder.createTextDataset(
       a,
-      e,
-      r,
+      s,
+      i,
       !0
     );
-    return { trainDataset: i, validationDataset: d };
+    return { trainDataset: o, validationDataset: n };
   }
-  async createDataset(t, e = 32) {
-    const s = await m(t, this.tokenizer);
-    return await this.datasetBuilder.createTextDataset(s, e);
+  async createDataset(t, s = 32) {
+    const e = await p(t, this.tokenizer);
+    return await this.datasetBuilder.createTextDataset(e, s);
   }
   dispose() {
     this.optimizer && this.optimizer.dispose();
   }
 }
 export {
-  R as default
+  x as default
 };

package/dist/training/sparseCrossEntropy.js

@@ -1,35 +1,28 @@
-import { gatherSub as L } from "../ops/gatherSub.js";
-import { scatterSub as y } from "../ops/scatterSub.js";
-import { e as u, l as i, z as S, t as f, b as G } from "../index-BAzbokzv.js";
-import { s as v } from "../softmax-fsdtf6JC.js";
-import { m as z } from "../max-DtlIuVeW.js";
-import { l as k } from "../log_sum_exp-YEo2h3gb.js";
-function F(a, s) {
-  return f(() => {
-    const e = a.shape[a.shape.length - 1], o = a.shape.slice(0, -1).reduce((d, c) => d * c, 1), p = a.shape.length > 2 ? a.reshape([o, e]) : a, n = s.shape.length > 1 ? s.reshape([o]).cast("int32") : s.cast("int32"), t = z(p, -1, !0), r = G(p, t), h = k(r, -1);
-    return L(h, n, r);
+import { gatherSub as x } from "../ops/gatherSub.js";
+import { scatterSub as L } from "../ops/scatterSub.js";
+import { q as C, t as u, z as E, c as G } from "../index-C0dhsYom.js";
+import { s as y } from "../softmax-GPNK3o-U.js";
+import { m as z } from "../max-BQc2Aj-I.js";
+import { l as v } from "../log_sum_exp-D086OgZJ.js";
+function k(t, s) {
+  return u(() => {
+    const n = t.shape[t.shape.length - 1], c = t.shape.slice(0, -1).reduce((o, e) => o * e, 1), h = t.shape.length > 2 ? t.reshape([c, n]) : t, p = s.shape.length > 1 ? s.reshape([c]).cast("int32") : s.cast("int32"), r = z(h, -1, !0), a = G(h, r), m = v(a, -1);
+    return x(m, p, a);
   });
 }
-function j() {
-  return u().backendName === "tensorflow" ? i((s, e, m) => {
-    const o = s.shape.length > 2 ? s.reshape([-1, s.shape[s.shape.length - 1]]) : s, p = e.shape.length > 1 ? e.reshape([-1]).cast("int32") : e.cast("int32"), [n, t] = u().runKernel(
-      "NativeSparseSoftmaxCrossEntropy",
-      { logits: o, labels: p },
-      {}
-    );
-    return m([t.reshape(s.shape)]), { value: n, gradFunc: (r, h) => [h[0], S(e)] };
-  }) : i(
+function A() {
+  return C(
     // @ts-expect-error Invalid params
-    (s, e, m) => {
-      const o = s.shape[s.shape.length - 1], n = s.shape.slice(0, -1).reduce((d, c) => d * c, 1), t = s.reshape([n, o]), r = e.reshape([n]).cast("int32"), h = F(t, r);
-      return m([t, r]), t.dispose(), r.dispose(), { value: h, gradFunc: (d, c) => f(() => {
-        const g = c[0], b = c[1], x = v(g), C = y(x, b, d), E = S(e);
-        return [C.reshape(s.shape), E];
+    (s, n, d) => {
+      const c = s.shape[s.shape.length - 1], p = s.shape.slice(0, -1).reduce((o, e) => o * e, 1), r = s.reshape([p, c]), a = n.reshape([p]).cast("int32"), m = k(r, a);
+      return d([r, a]), r.dispose(), a.dispose(), { value: m, gradFunc: (o, e) => u(() => {
+        const S = e[0], f = e[1], b = y(S), l = L(b, f, o), g = E(n);
+        return [l.reshape(s.shape), g];
       }) };
     }
   );
 }
 export {
-  j as createSoftmaxCrossEntropyWithGrad,
-  F as sparseSoftmaxCrossEntropy
+  A as createSoftmaxCrossEntropyWithGrad,
+  k as sparseSoftmaxCrossEntropy
 };

package/dist/utilities/dummy.js

@@ -1,35 +1,35 @@
-import { m as f, v as S, e as w } from "../index-BAzbokzv.js";
-import { z as i } from "../zeros-YJDE7oRb.js";
-async function P(s) {
+import { m as y, v as P, e as S } from "../index-C0dhsYom.js";
+import { z as i } from "../zeros-DCPCdFGq.js";
+async function w(s) {
   const t = i([1, s.config.gpt.blockSize], "int32"), [e, n] = s.forward({ training: !1 }, t);
   await e.data(), e.dispose(), n && n.dispose(), t.dispose();
 }
-async function v(s) {
-  const t = f(), e = t.numBytesInGPUAllocated ?? t.numBytes;
-  await P(s);
+async function k(s) {
+  const t = y(), e = t.numBytesInGPUAllocated ?? t.numBytesAllocatedInGPU ?? t.numBytes;
+  await w(s);
   const n = i([1, s.config.gpt.blockSize], "int32"), r = i([1, s.config.gpt.blockSize], "int32"), o = {
     perBatch: 0,
     tapeSize: 0,
     gradients: s.getNumParams() * 4
-  }, g = () => {
-    const [a, l] = s.forward({ training: !0 }, n, r), d = w().state.activeTape;
+  }, f = () => {
+    const [c, l] = s.forward({ training: !0 }, n, r), p = S().state.activeTape;
     let u = 0;
-    if (d)
-      for (const z of d)
+    if (p)
+      for (const z of p)
         u += z.saved?.reduce((B, I) => B + I.size * 4, 0) || 0;
-    return o.tapeSize = u, a.dispose(), l;
-  }, { value: c, grads: m } = S(g), p = f(), y = p.numBytesInGPUAllocated ?? p.numBytes;
-  o.perBatch = y - e - o.gradients, console.log("Dummy training memory requirements:", o), await c.data(), c.dispose();
-  for (const a in m)
-    m[a].dispose();
+    return o.tapeSize = u, c.dispose(), l;
+  }, { value: m, grads: d } = P(f), a = y(), g = a.numBytesInGPUAllocated ?? a.numBytesAllocatedInGPU ?? a.numBytes;
+  o.perBatch = g - e - o.gradients, console.log("Dummy training memory requirements:", o), await m.data(), m.dispose();
+  for (const c in d)
+    d[c].dispose();
   return n.dispose(), r.dispose(), o;
 }
-function A(s) {
+function v(s) {
   const t = i([1, s.config.gpt.blockSize], "int32"), [e, n] = s.forward({ training: !1 }, t);
   e.dispose(), n && n.dispose(), t.dispose();
 }
 export {
-  A as dummyPass,
-  P as dummyPassAsync,
-  v as dummyPassTrainAsync
+  v as dummyPass,
+  w as dummyPassAsync,
+  k as dummyPassTrainAsync
 };

package/dist/utilities/generate.js

@@ -1,23 +1,22 @@
-import { t as y } from "../index-BAzbokzv.js";
-import { t as x } from "../tensor2d-CPEkynbH.js";
-import { c as f } from "../concat-5aPGqw3Z.js";
-async function A(o, r, a, c, T) {
+import "../index-C0dhsYom.js";
+import { t as m } from "../tensor2d-B_kyod7_.js";
+import { c as u } from "../concat-CsHeR4zV.js";
+async function v(o, r, a, c, f) {
   if (c <= 0)
     throw new Error("Length must be a positive integer");
   if (a.length === 0)
     throw new Error("Prompt cannot be an empty string");
-  const p = await o.tokenise([a], !0), s = r.config.gpt.useRope ? new Array(r.config.gpt.nLayer).fill(void 0) : void 0, d = y(() => {
-    let t = x(p, [1, p[0].length], "int32"), e = t;
-    for (let g = 0; g < c; g++) {
-      const { output: n } = r.generate(t, s, T), l = t, m = e;
-      e = f([e, n], 1), t = s ? n : f([t, n], 1), l.dispose(), m.dispose(), s || n.dispose();
-    }
-    return e;
-  }), w = await d.array();
-  d.dispose();
-  const i = w[0], u = i.indexOf(o.eosToken);
-  return u !== -1 && i.splice(u), await o.decode(i);
+  const p = await o.tokenise([a], !0), s = r.config.gpt.useRope ? new Array(r.config.gpt.nLayer).fill(void 0) : void 0;
+  let e = m(p, [1, p[0].length], "int32"), t = e;
+  for (let g = 0; g < c; g++) {
+    const { output: n } = await r.generate(e, s, f), T = e, l = t;
+    t = u([t, n], 1), e = s ? n : u([e, n], 1), T.dispose(), l.dispose(), s || n.dispose();
+  }
+  const w = await t.array();
+  t.dispose();
+  const i = w[0], d = i.indexOf(o.eosToken);
+  return d !== -1 && i.splice(d), await o.decode(i);
 }
 export {
-  A as generateText
+  v as generateText
 };

package/dist/utilities/multinomialCPU.d.ts

@@ -0,0 +1,2 @@
+import { Tensor2D } from '@tensorflow/tfjs-core';
+export default function multinomialCPU(probs: number[]): Tensor2D;

package/dist/utilities/multinomialCPU.js

@@ -0,0 +1,13 @@
+import "../index-C0dhsYom.js";
+import { t as e } from "../tensor2d-B_kyod7_.js";
+function l(n) {
+  let r = 0;
+  const i = Math.random();
+  for (let t = 0; t < n.length; t++)
+    if (r += n[t], i < r)
+      return e([[t]], [1, 1], "int32");
+  return e([[n.length - 1]], [1, 1], "int32");
+}
+export {
+  l as default
+};

package/dist/utilities/performance.d.ts

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export default function performanceTest(fn: () => Tensor, iterations?: number): Promise<number>;

package/dist/utilities/performance.js

@@ -0,0 +1,16 @@
+import { t as r } from "../index-C0dhsYom.js";
+async function d(s, o = 10) {
+  for (let e = 0; e < 10; e++) {
+    const t = s();
+    await t.data(), t.dispose();
+  }
+  const a = performance.now();
+  for (let e = 0; e < o; e++) {
+    const t = r(s);
+    e === o - 1 && await t.data(), t.dispose();
+  }
+  return (performance.now() - a) / o;
+}
+export {
+  d as default
+};

package/dist/utilities/profile.d.ts

@@ -1,6 +1,7 @@
 import { MemoryInfo } from '@tensorflow/tfjs-core';
 export interface ExtendedMemoryInfo extends MemoryInfo {
     numBytesInGPUAllocated?: number;
+    numBytesAllocatedInGPU?: number;
 }
 export default class MemoryProfiler {
     private log;