@genai-fi/nanogpt 0.9.1 → 0.10.1

package/dist/training/Adam.js

@@ -1,26 +1,10 @@
 import { adamAdjust as b } from "../ops/adamAdjust.js";
 import { adamMoments as d } from "../ops/adamMoments.js";
-import { O as g, e as h, t as o, d as B } from "../index-BzFyqcy-.js";
-import { z as M } from "../zeros-Bj5rMYA7.js";
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
+import { O as g, e as h, t as o, d as B } from "../index-ZyQhjEPo.js";
+import { z as M } from "../zeros-2gldETuK.js";
 class R extends g {
-  constructor(t, a, e, s = null) {
-    super(), this.learningRate = t, this.beta1 = a, this.beta2 = e, this.epsilon = s, this.accBeta1 = a, this.accBeta2 = e, s === null && (this.epsilon = h().backend.epsilon());
+  constructor(t, a, e, s, i = null) {
+    super(), this.learningRate = t, this.beta1 = a, this.beta2 = e, this.lossScaling = s, this.epsilon = i, this.accBeta1 = a, this.accBeta2 = e, i === null && (this.epsilon = h().backend.epsilon());
   }
   /** @nocollapse */
   static get className() {
@@ -33,19 +17,19 @@ class R extends g {
     const a = Array.isArray(t) ? t.map((e) => e.name) : Object.keys(t);
     o(() => {
       const e = 1 - this.accBeta1, s = 1 - this.accBeta2;
-      a.forEach((n, i) => {
-        const c = h().registeredVariables[n], u = !1;
-        this.accumulatedMoments[i] == null && (this.accumulatedMoments[i] = {
-          originalName: `${n}/m`,
+      a.forEach((i, n) => {
+        const c = h().registeredVariables[i], u = !1;
+        this.accumulatedMoments[n] == null && (this.accumulatedMoments[n] = {
+          originalName: `${i}/m`,
           variable: o(() => M([...c.shape, 2]).variable(u))
         });
-        const r = Array.isArray(t) ? t[i].tensor : t[n];
+        const r = Array.isArray(t) ? t[n].tensor : t[i];
         if (r == null)
           return;
-        const m = this.accumulatedMoments[i].variable, l = d(m, r, this.beta1, this.beta2);
-        m.assign(l);
+        const l = this.accumulatedMoments[n].variable, m = d(l, r, this.beta1, this.beta2, this.lossScaling);
+        l.assign(m);
         const p = b(
-          l,
+          m,
           c,
           e,
           s,

package/dist/training/AdamExt.d.ts

@@ -5,6 +5,7 @@ interface AdamExtConfig {
     decaySteps: number;
     minLearningRate: number;
     weightDecay?: number;
+    lossScaling: number;
 }
 /**
  * Extended Adam optimizer with warmup, cosine decay, and optional weight decay.

package/dist/training/AdamExt.js

@@ -1,8 +1,8 @@
-import { a as r, b as c, c as h, e as o } from "../index-BzFyqcy-.js";
+import { m as r, b as c, c as h, e as o } from "../index-ZyQhjEPo.js";
 import { AdamOptimizer as g } from "./Adam.js";
 class y extends g {
   constructor(t, e, s, i, a) {
-    super(t, e, s, i), this.config = a, this.startLearningRate = t;
+    super(t, e, s, a.lossScaling, i), this.config = a, this.startLearningRate = t;
   }
   step = 0;
   startLearningRate;

package/dist/training/DatasetBuilder.js

@@ -1,23 +1,6 @@
-import { t as g } from "../index-BzFyqcy-.js";
-import { d as u, i as d } from "../dataset-DlZtKmBq.js";
-import "../index-Tf7vU29b.js";
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- *
- * =============================================================================
- */
+import { t as g } from "../index-ZyQhjEPo.js";
+import { d as u, i as d } from "../dataset-0xP8GjwI.js";
+import "../index-Cp39cXWe.js";
 function z(r) {
   return u(async () => {
     const t = await r();

package/dist/training/FullTrainer.js

@@ -1,13 +1,14 @@
-import L from "./Trainer.js";
-import w from "./Evaluator.js";
-import { d as S } from "../index-BzFyqcy-.js";
-import f from "../utilities/profile.js";
-const y = {
+import b from "./Trainer.js";
+import L from "./Evaluator.js";
+import { d as w } from "../index-ZyQhjEPo.js";
+import y from "../utilities/profile.js";
+import { createTensorStatistics as D } from "../checks/weights.js";
+const T = {
   desiredLoss: 0.01,
   logInterval: 1,
   maxSteps: 1e3
 };
-class x extends L {
+class V extends b {
   constructor(s, e, a = 3e-4) {
     super(s, e, a);
   }
@@ -42,89 +43,95 @@ class x extends L {
   }
   async stepDataset(s, e, a) {
     const { logInterval: l } = {
-      ...y,
+      ...T,
       ...e
     }, c = Date.now(), r = this.createEmptyState();
-    this.lastState = r, await this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || this.model.setProfiler(new f())), this.running = !0, r.logStartTime = c;
-    const d = a ? new w(this.model, a) : void 0, t = await s.iterator();
+    this.lastState = r, await this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || this.model.setProfiler(new y())), this.running = !0, r.logStartTime = c;
+    const g = a ? new L(this.model, a) : void 0, t = await s.iterator();
     try {
       for (; this.running; ) {
-        const i = await t.next();
-        if (i.done) break;
-        const m = i.value, o = this.trainBatch(r, m);
+        const n = await t.next();
+        if (n.done) break;
+        const m = n.value, o = this.trainBatch(r, m);
         if (r.step % l === 0) {
-          const g = (await o.data())[0];
-          r.lastLoss = g;
+          const p = (await o.data())[0];
+          r.lastLoss = p;
           const u = Date.now();
           r.trainingDuration += u - r.logStartTime;
-          const p = this.createLogEntry(r, c, m.xs.shape[0], e?.advancedMetrics);
+          const d = this.createLogEntry(r, c, m.xs.shape[0], e?.advancedMetrics);
           if (this.model.trainingState = {
             steps: r.totalSteps,
             learningRate: this.optimizer.lr,
             batchSize: m.xs.shape[0],
             loss: r.lastLoss
-          }, d)
+          }, g)
             try {
-              const n = await d.evaluate(5);
-              r.validationLosses.push(n), p.valLoss = n;
-            } catch (n) {
-              console.error("Validation error:", n);
+              const h = await g.evaluate(5);
+              r.validationLosses.push(h), d.valLoss = h;
+            } catch (h) {
+              console.error("Validation error:", h);
             }
-          const v = this.createProgress(r, p, e?.advancedMetrics);
-          return o.dispose(), this.stop(), { log: p, progress: v };
+          const S = this.createProgress(r, d, e?.advancedMetrics);
+          return o.dispose(), this.stop(), { log: d, progress: S };
         }
         o.dispose();
       }
-    } catch (i) {
-      throw console.error("Training error:", i), S(), i;
+    } catch (n) {
+      throw console.error("Training error:", n), w(), n;
     }
-    throw S(), this.running = !1, new Error("No log returned before training stopped.");
+    throw w(), this.running = !1, new Error("No log returned before training stopped.");
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
   async trainOnDataset(s, e, a) {
     const { logInterval: l, onStep: c, maxSteps: r } = {
-      ...y,
+      ...T,
       ...e
-    }, d = Date.now(), t = this.createEmptyState();
-    this.lastState = t, await this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || this.model.setProfiler(new f())), this.running = !0, t.logStartTime = d;
-    const i = a ? new w(this.model, a) : void 0, m = await s.iterator();
+    }, g = Date.now(), t = this.createEmptyState();
+    this.lastState = t, await this.dummyPass(), this.model.trainable = !0, e?.advancedMetrics && (this.model.getProfiler() || this.model.setProfiler(new y())), console.log("Training options", e), this.running = !0, t.logStartTime = g;
+    const n = a ? new L(this.model, a) : void 0, m = await s.iterator();
     try {
       for (; this.running; ) {
         const o = await m.next();
         if (o.done) break;
-        const g = o.value, u = this.trainBatch(t, g);
-        if (t.step % l === 0) {
-          const p = (await u.data())[0];
-          t.lastLoss = p;
-          const v = Date.now();
-          t.trainingDuration += v - t.logStartTime;
-          const n = this.createLogEntry(t, d, g.xs.shape[0], e?.advancedMetrics);
+        const p = o.value, u = t.step % l === 0, d = (e?.gradientMetrics || !1) && u, S = this.trainBatch(t, p, d);
+        if (u) {
+          const h = (await S.data())[0];
+          t.lastLoss = h;
+          const P = Date.now();
+          t.trainingDuration += P - t.logStartTime;
+          const f = this.createLogEntry(t, g, p.xs.shape[0], e?.advancedMetrics);
           if (this.model.trainingState = {
             steps: t.totalSteps,
             learningRate: this.optimizer.lr,
-            batchSize: g.xs.shape[0],
+            batchSize: p.xs.shape[0],
             loss: t.lastLoss
-          }, i)
+          }, e?.gradientMetrics && d && t.gradients) {
+            const i = /* @__PURE__ */ new Map();
+            for (const [M, v] of Object.entries(t.gradients))
+              i.set(M, await D(v)), v.dispose();
+            f.gradientMetrics = i;
+          }
+          if (n)
             try {
-              const h = await i.evaluate(5);
-              t.validationLosses.push(h), n.valLoss = h;
-            } catch (h) {
-              console.error("Validation error:", h);
+              const i = await n.evaluate(5);
+              t.validationLosses.push(i), f.valLoss = i;
+            } catch (i) {
+              console.error("Validation error:", i);
             }
           if (c) {
-            const h = this.createProgress(t, n, e?.advancedMetrics);
-            await c(n, h);
+            const i = this.createProgress(t, f, e?.advancedMetrics);
+            await c(f, i);
           }
           t.logStartTime = Date.now();
         }
-        u.dispose(), t.step >= r && this.stop();
+        S.dispose(), t.step >= r && this.stop();
       }
     } catch (o) {
-      throw console.error("Training error:", o), S(), o;
+      throw console.error("Training error:", o), w(), o;
     }
-    return S(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
+    return w(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  x as default
+  V as default
 };

package/dist/training/Trainer.d.ts

@@ -1,10 +1,11 @@
 import { ITokeniser } from '../tokeniser/type';
 import { DatasetBuilder } from './DatasetBuilder';
 import { default as AdamExt } from './AdamExt';
-import { TensorContainer } from '@tensorflow/tfjs-core/dist/tensor_types';
+import { NamedTensorMap, TensorContainer } from '@tensorflow/tfjs-core/dist/tensor_types';
 import { Scalar, Tensor } from '@tensorflow/tfjs-core';
 import { Dataset } from '@tensorflow/tfjs-data';
 import { default as Model, ModelForwardAttributes } from '../models/model';
+import { TensorStatistics } from '../checks/weights';
 export interface TrainingLogEntry {
     loss: number;
     valLoss?: number;
@@ -12,7 +13,7 @@ export interface TrainingLogEntry {
     time: number;
     example?: string;
     batchSize: number;
-    gradientNorm?: number;
+    gradientMetrics?: Map<string, TensorStatistics>;
     learningRate?: number;
 }
 export interface TrainingState {
@@ -21,7 +22,7 @@ export interface TrainingState {
     totalSteps: number;
     losses: number[];
     validationLosses: number[];
-    gradientNorm?: Promise<number>;
+    gradients?: NamedTensorMap;
 }
 export interface TrainingProgress {
     duration: number;
@@ -41,6 +42,7 @@ export interface TrainingOptions {
     prompt?: string;
     maxSteps: number;
     advancedMetrics?: boolean;
+    gradientMetrics?: boolean;
     onStep?: (log: TrainingLogEntry, progress: TrainingProgress) => Promise<void> | void;
 }
 export default abstract class GPTTrainer {
@@ -52,22 +54,25 @@ export default abstract class GPTTrainer {
     protected running: boolean;
     protected lastState?: TrainingState;
     protected _gradientCheckpointing: boolean;
+    protected _mixedPrecision: boolean;
+    protected lossScaling: number;
     constructor(model: Model<ModelForwardAttributes>, tokenizer: ITokeniser, learningRate?: number);
     setGradientCheckpointing(enabled: boolean): void;
+    setMixedPrecision(enabled: boolean): void;
     setLearningRate(learningRate: number): void;
     reset(): void;
     stop(): void;
     getOptimizer(): AdamExt;
     resetOptimizer(config?: AdamConfig): void;
-    protected trainStep(_state: Partial<TrainingState>, batch: {
+    protected trainStep(state: Partial<TrainingState>, batch: {
         xs: Tensor;
         ys: Tensor;
-    }, dummy?: boolean): Scalar;
+    }, dummy?: boolean, keepGrads?: boolean): Scalar;
     protected dummyPass(): Promise<void>;
     protected trainBatch(state: TrainingState, batch: {
         xs: Tensor;
         ys: Tensor;
-    }): Scalar;
+    }, keepGrads?: boolean): Scalar;
     abstract trainOnDataset(dataset: Dataset<{
         xs: Tensor;
         ys: Tensor;

package/dist/training/Trainer.js

@@ -1,10 +1,10 @@
-import { DatasetBuilder as m, flattenTokens as c, PAGE_FACTOR as g } from "./DatasetBuilder.js";
-import u from "./AdamExt.js";
-import { t as f, v as y, d as p } from "../index-BzFyqcy-.js";
-import { z as h } from "../zeros-Bj5rMYA7.js";
-class x {
-  constructor(t, e, i = 1e-3) {
-    this.tokenizer = e, this.model = t, this.learningRate = i, this.resetOptimizer(), this.datasetBuilder = new m(e, t.config.blockSize);
+import { DatasetBuilder as f, flattenTokens as h, PAGE_FACTOR as y } from "./DatasetBuilder.js";
+import z from "./AdamExt.js";
+import { t as S, v as k, k as x, d as p, b as m } from "../index-ZyQhjEPo.js";
+import { z as g } from "../zeros-2gldETuK.js";
+class M {
+  constructor(t, e, s = 1e-3) {
+    this.tokenizer = e, this.model = t, this.lossScaling = t.lossScaling, this.learningRate = s, this.resetOptimizer(), this.datasetBuilder = new f(e, t.config.blockSize);
   }
   model;
   optimizer;
@@ -13,9 +13,14 @@ class x {
   running = !1;
   lastState;
   _gradientCheckpointing = !1;
+  _mixedPrecision = !1;
+  lossScaling;
   setGradientCheckpointing(t) {
     this._gradientCheckpointing = t;
   }
+  setMixedPrecision(t) {
+    this._mixedPrecision = t;
+  }
   setLearningRate(t) {
     this.learningRate = t, this.resetOptimizer({ learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 });
   }
@@ -30,7 +35,7 @@ class x {
   }
   resetOptimizer(t = { learningRateFactor: 1, beta1: 0.9, beta2: 0.99, epsilon: 1e-8 }) {
     this.optimizer && this.optimizer.dispose();
-    const e = new u(
+    const e = new z(
       t.learningRateFactor * this.learningRate,
       t.beta1,
       t.beta2,
@@ -39,69 +44,76 @@ class x {
         warmupSteps: 100,
         decaySteps: 2e4,
         minLearningRate: 1e-4,
-        weightDecay: 0
+        weightDecay: 0,
+        lossScaling: this.lossScaling
       }
     );
     this.optimizer = e;
   }
-  trainStep(t, e, i = !1) {
-    return f(() => {
+  trainStep(t, e, s = !1, i = !1) {
+    return S(() => {
       this.model.getProfiler()?.startMemory();
-      const { xs: a, ys: s } = e, n = () => {
-        const [l, d] = this.model.forward(
-          { training: !0, checkpointing: this._gradientCheckpointing },
+      const { xs: a, ys: l } = e, c = () => {
+        const [n, d] = this.model.forward(
+          {
+            training: !0,
+            checkpointing: this._gradientCheckpointing,
+            mixedPrecision: this._mixedPrecision
+          },
           a,
-          s
+          l
         );
-        return l.dispose(), d;
-      }, { value: o, grads: r } = y(n);
-      return i ? this.model.getProfiler()?.endMemory("Training") : (this.optimizer.applyGradients(r), this.model.getProfiler()?.endMemory("Training"), p(r)), o;
+        n.dispose();
+        const u = d.mul(m(this.lossScaling));
+        return d.dispose(), u;
+      }, { value: o, grads: r } = k(c);
+      return s ? this.model.getProfiler()?.endMemory("Training") : (this.optimizer.applyGradients(r), this.model.getProfiler()?.endMemory("Training"), i ? (t.gradients = r, Object.values(r).forEach((n) => x(n))) : p(r)), o.mul(m(1 / this.lossScaling));
     });
   }
   async dummyPass() {
-    const t = h([1, this.model.config.blockSize], "int32"), e = h([1, this.model.config.blockSize], "int32");
+    const t = g([1, this.model.config.blockSize], "int32"), e = g([1, this.model.config.blockSize], "int32");
     try {
-      const i = this.trainStep({}, { xs: t, ys: e }, !0);
-      await i.data(), i.dispose();
-    } catch (i) {
-      console.error("Error during dummy pass:", i);
+      const s = this.trainStep({}, { xs: t, ys: e }, !0);
+      await s.data(), s.dispose();
+    } catch (s) {
+      console.error("Error during dummy pass:", s);
     } finally {
       t.dispose(), e.dispose();
     }
   }
-  trainBatch(t, e) {
+  trainBatch(t, e, s = !1) {
     try {
-      const i = this.trainStep(t, e, !1);
+      const i = this.trainStep(t, e, !1, s);
       return e.xs.dispose(), e.ys.dispose(), t.step++, t.totalSteps++, i;
     } catch (i) {
       throw console.error(`Error processing batch at step ${t.step}:`, i), p(), i;
     }
   }
-  async createTrainValidationSplit(t, e = 32, i = 0.1) {
-    const a = await c(t, this.tokenizer), s = /* @__PURE__ */ new Set();
-    if (i > 0) {
-      const r = Math.floor(a.length / (this.datasetBuilder.blockSize * g)), l = Math.max(1, Math.floor(r * i));
-      for (; s.size < l; ) {
-        const d = Math.floor(Math.random() * r);
-        s.add(d);
+  async createTrainValidationSplit(t, e = 32, s = 0.1) {
+    const i = await h(t, this.tokenizer), a = /* @__PURE__ */ new Set();
+    if (s > 0) {
+      const o = Math.floor(i.length / (this.datasetBuilder.blockSize * y)), r = Math.max(1, Math.floor(o * s));
+      for (; a.size < r; ) {
+        const n = Math.floor(Math.random() * o);
+        a.add(n);
       }
     }
-    const n = await this.datasetBuilder.createTextDataset(a, e, s, !1), o = await this.datasetBuilder.createTextDataset(
-      a,
+    const l = await this.datasetBuilder.createTextDataset(i, e, a, !1), c = await this.datasetBuilder.createTextDataset(
+      i,
       e,
-      s,
+      a,
       !0
     );
-    return { trainDataset: n, validationDataset: o };
+    return { trainDataset: l, validationDataset: c };
   }
   async createDataset(t, e = 32) {
-    const i = await c(t, this.tokenizer);
-    return await this.datasetBuilder.createTextDataset(i, e);
+    const s = await h(t, this.tokenizer);
+    return await this.datasetBuilder.createTextDataset(s, e);
   }
   dispose() {
     this.optimizer && this.optimizer.dispose();
   }
 }
 export {
-  x as default
+  M as default
 };

package/dist/training/sparseCrossEntropy.js

@@ -1,8 +1,8 @@
 import { gatherSub as x } from "../ops/gatherSub.js";
 import { scatterSub as L } from "../ops/scatterSub.js";
-import { z as C, t as u, A as E, c as G } from "../index-BzFyqcy-.js";
-import { s as y } from "../softmax-D7Jj3p_P.js";
-import { m as z, l as v } from "../log_sum_exp-DO6z8tSE.js";
+import { w as C, t as u, z as E, c as G } from "../index-ZyQhjEPo.js";
+import { s as y } from "../softmax-ZHVebtR1.js";
+import { m as z, l as v } from "../log_sum_exp-DWI-76TI.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), d = v(a, -1);

package/dist/transpose-DKELTqhe.js

@@ -0,0 +1,38 @@
+import { A as u, B as i, E as a, t as m } from "./index-ZyQhjEPo.js";
+import { D as $, N as g, F as x, H as c } from "./tensor_util-DV-FP5Q3.js";
+import { c as k } from "./complex-CSlYz-2T.js";
+import { a as l } from "./tensor-DdQUJZlz.js";
+function K(r) {
+  const e = { input: i(r, "input", "imag") };
+  return a.runKernel($, e);
+}
+const h = /* @__PURE__ */ u({ imag_: K });
+function E(r) {
+  const e = { x: i(r, "x", "neg") };
+  return a.runKernel(g, e);
+}
+const N = /* @__PURE__ */ u({ neg_: E });
+function _(r) {
+  const e = { input: i(r, "input", "real") };
+  return a.runKernel(x, e);
+}
+const b = /* @__PURE__ */ u({ real_: _ });
+function d(r, t, e) {
+  const n = i(r, "x", "transpose");
+  if (t == null && (t = n.shape.map((s, o) => o).reverse()), l(n.rank === t.length, () => `Error in transpose: rank of input ${n.rank} must match length of perm ${t}.`), t.forEach((s) => {
+    l(s >= 0 && s < n.rank, () => `All entries in 'perm' must be between 0 and ${n.rank - 1} but got ${t}`);
+  }), n.rank <= 1)
+    return n.clone();
+  const f = { x: n }, p = { perm: t };
+  return n.dtype === "complex64" ? m(() => {
+    let s = b(n), o = h(n);
+    return s = a.runKernel(c, { x: s }, p), o = a.runKernel(c, { x: o }, p), e && (o = N(o)), k(s, o);
+  }) : a.runKernel(c, f, p);
+}
+const I = /* @__PURE__ */ u({ transpose_: d });
+export {
+  h as i,
+  N as n,
+  b as r,
+  I as t
+};

package/dist/utilities/arrayClose.js

@@ -1,20 +1,20 @@
-function n(r, e) {
+function i(r, e) {
   let t = 0;
-  if (Array.isArray(r) && Array.isArray(e)) {
+  if ((Array.isArray(r) || r instanceof Float32Array) && (Array.isArray(e) || e instanceof Float32Array)) {
     if (r.length !== e.length) return Number.POSITIVE_INFINITY;
-    for (let i = 0; i < r.length; ++i)
-      t = Math.max(t, n(r[i], e[i]));
+    for (let n = 0; n < r.length; ++n)
+      t = Math.max(t, i(r[n], e[n]));
     return t;
   } else if (typeof r == "number" && typeof e == "number") {
     if (isNaN(r) && isNaN(e))
       return 0;
     if (!isFinite(r) || !isFinite(e))
       return r === e ? 0 : Number.POSITIVE_INFINITY;
-    const i = Math.abs(r - e);
-    return t = Math.max(t, i), t;
+    const n = Math.abs(r - e);
+    return t = Math.max(t, n), t;
   } else
     return Number.POSITIVE_INFINITY;
 }
 export {
-  n as arraysClose
+  i as arraysClose
 };

package/dist/utilities/dummy.js

@@ -1,35 +1,43 @@
-import { m as y, v as P, e as S } from "../index-BzFyqcy-.js";
-import { z as i } from "../zeros-Bj5rMYA7.js";
-async function w(s) {
-  const t = i([1, s.config.blockSize], "int32"), [e, n] = s.forward({ training: !1 }, t);
-  await e.data(), e.dispose(), n && n.dispose(), t.dispose();
+import { a as y, e as S, v as w } from "../index-ZyQhjEPo.js";
+import { z as m } from "../zeros-2gldETuK.js";
+import { o as P } from "../ones-CAMiP4I2.js";
+async function b(s) {
+  const t = m([1, s.config.blockSize], "int32"), [n, o] = s.forward({ training: !1 }, t);
+  await n.data(), n.dispose(), o && o.dispose(), t.dispose();
 }
-async function k(s) {
-  const t = y(), e = t.numBytesInGPUAllocated ?? t.numBytesAllocatedInGPU ?? t.numBytes;
-  await w(s);
-  const n = i([1, s.config.blockSize], "int32"), r = i([1, s.config.blockSize], "int32"), o = {
+async function G(s) {
+  console.log("Starting dummy training pass for memory profiling...");
+  const t = y(), n = t.numBytesInGPUAllocated ?? t.numBytesAllocatedInGPU ?? t.numBytes;
+  await b(s), console.log("Forward pass complete. Starting backward pass...");
+  const o = m([1, s.config.blockSize], "int32"), a = P([1, s.config.blockSize], "int32"), e = {
     perBatch: 0,
     tapeSize: 0,
     gradients: s.getNumParams() * 4
-  }, f = () => {
-    const [c, g] = s.forward({ training: !0 }, n, r), u = S().state.activeTape;
-    let p = 0;
-    if (u)
-      for (const z of u)
-        p += z.saved?.reduce((B, I) => B + I.size * 4, 0) || 0;
-    return o.tapeSize = p, c.dispose(), g;
-  }, { value: m, grads: d } = P(f), a = y(), l = a.numBytesInGPUAllocated ?? a.numBytesAllocatedInGPU ?? a.numBytes;
-  o.perBatch = l - 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;
+  };
+  try {
+    const i = () => {
+      const [c, g] = s.forward({ training: !0 }, o, a), u = S().state.activeTape;
+      let l = 0;
+      if (u)
+        for (const z of u)
+          l += z.saved?.reduce((B, I) => B + I.size * 4, 0) || 0;
+      return e.tapeSize = l, c.dispose(), g;
+    }, { value: d, grads: p } = w(i), r = y(), f = r.numBytesInGPUAllocated ?? r.numBytesAllocatedInGPU ?? r.numBytes;
+    e.perBatch = f - n - e.gradients, console.log("Dummy training memory requirements:", e), await d.data(), d.dispose();
+    for (const c in p)
+      p[c].dispose();
+    o.dispose(), a.dispose();
+  } catch (i) {
+    console.error("Error during dummy training pass:", i), o.dispose(), a.dispose();
+  }
+  return e;
 }
-function v(s) {
-  const t = i([1, s.config.blockSize], "int32"), [e, n] = s.forward({ training: !1 }, t);
-  e.dispose(), n && n.dispose(), t.dispose();
+function T(s) {
+  const t = m([1, s.config.blockSize], "int32"), [n, o] = s.forward({ training: !1 }, t);
+  n.dispose(), o && o.dispose(), t.dispose();
 }
 export {
-  v as dummyPass,
-  w as dummyPassAsync,
-  k as dummyPassTrainAsync
+  T as dummyPass,
+  b as dummyPassAsync,
+  G as dummyPassTrainAsync
 };

package/dist/utilities/multinomialCPU.js

@@ -1,5 +1,5 @@
-import "../index-BzFyqcy-.js";
-import { t as e } from "../tensor2d-D76QGjF3.js";
+import "../index-ZyQhjEPo.js";
+import { t as e } from "../tensor2d-G4Ys2GxX.js";
 function l(n) {
   let r = 0;
   const i = Math.random();