@genai-fi/nanogpt 0.13.0 → 0.13.1

package/dist/TeachableLLM.js

@@ -36,10 +36,10 @@ import c from "./tokeniser/CharTokeniser.js";
 import g from "./tokeniser/bpe.js";
 import "./papaparse.min-C0cScC2i.js";
 import "./jszip.min-Bz5-11Bk.js";
-import "./ops/cpu/scatterSub.js";
-import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/gatherSub.js";
 import "./ops/webgl/gatherSub.js";
+import "./ops/cpu/scatterSub.js";
+import "./ops/webgl/scatterSub.js";
 import "./ops/cpu/matMulGelu.js";
 import "./matMulGelu-CoUYwB2k.js";
 import "./ops/grads/matMulGelu.js";

package/dist/main.d.ts

@@ -14,7 +14,9 @@ export { default as BPETokeniser } from './tokeniser/bpe';
 export { default as waitForModel } from './utilities/waitForModel';
 export { default as loadTextData } from './data/textLoader';
 export { default as Generator } from './Generator';
+export { default as Evaluator } from './training/Evaluator';
 export type { ITrainerOptions } from './Trainer';
+export { default as Trainer } from './Trainer';
 export type { IGenerateOptions } from './Generator';
 export { type ModelForwardAttributes, default as Model } from './models/model';
 export type { ITokeniser, Conversation, Roles } from './tokeniser/type';

package/dist/main.js

@@ -6,14 +6,16 @@ import "./index-Cp39cXWe.js";
 import "./dataset-BMe3pbsL.js";
 import { default as fo } from "./models/NanoGPTV1.js";
 import { default as lo } from "./TeachableLLM.js";
-import { default as co } from "./tokeniser/CharTokeniser.js";
+import { default as uo } from "./tokeniser/CharTokeniser.js";
 import { default as ko } from "./tokeniser/bpe.js";
 import { default as go } from "./utilities/waitForModel.js";
 import { default as Co } from "./data/textLoader.js";
-import { default as Eo } from "./models/model.js";
-import { estimateMemoryUsage as Bo, estimateParameterCount as bo, estimateResources as yo, estimateTrainingMemoryUsage as Lo, validateConfig as So } from "./utilities/parameters.js";
-import { default as Ao } from "./utilities/topP.js";
-import { Task as Go, tokensFromTasks as Ro } from "./training/tasks/Task.js";
+import { default as Mo } from "./training/Evaluator.js";
+import { default as Bo } from "./Trainer.js";
+import { default as vo } from "./models/model.js";
+import { estimateMemoryUsage as Lo, estimateParameterCount as So, estimateResources as Ao, estimateTrainingMemoryUsage as Fo, validateConfig as Go } from "./utilities/parameters.js";
+import { default as Uo } from "./utilities/topP.js";
+import { Task as Do, tokensFromTasks as No } from "./training/tasks/Task.js";
 import o from "./training/tasks/PretrainingTask.js";
 import r from "./training/tasks/StartSentenceTask.js";
 import t from "./training/tasks/ConversationTask.js";
@@ -52,15 +54,15 @@ import "./matMul16-CH8D42Kx.js";
 import "./ops/webgl/matMul16.js";
 import "./ops/cpu/matMul16.js";
 import "./ops/transpose16.js";
-import { selectBackend as wo } from "./backend.js";
-import { default as No } from "./utilities/performance.js";
-import p from "./layers/CausalSelfAttention.js";
-import a from "./layers/MLP.js";
+import { selectBackend as qo } from "./backend.js";
+import { default as Ho } from "./utilities/performance.js";
+import a from "./layers/CausalSelfAttention.js";
+import p from "./layers/MLP.js";
 import i from "./layers/TransformerBlock.js";
 import s from "./layers/RoPECache.js";
-import { default as qo } from "./training/AdamExt.js";
-import { default as Ho } from "./checks/index.js";
-import { sentenceEmbeddings as Jo, sentenceEmbeddingsTensor as Ko } from "./utilities/sentences.js";
+import { default as Jo } from "./training/AdamExt.js";
+import { default as Oo } from "./checks/index.js";
+import { sentenceEmbeddings as Vo, sentenceEmbeddingsTensor as Wo } from "./utilities/sentences.js";
 const to = {
   PretrainingTask: o,
   StartSentenceTask: r,
@@ -69,35 +71,37 @@ const to = {
   pack16: m,
   unpack16: e
 }, mo = {
-  CausalSelfAttention: p,
-  MLP: a,
+  CausalSelfAttention: a,
+  MLP: p,
   TransformerBlock: i,
   RoPECache: s
 };
 export {
-  qo as AdamExt,
+  Jo as AdamExt,
   ko as BPETokeniser,
-  co as CharTokeniser,
+  uo as CharTokeniser,
+  Mo as Evaluator,
   io as Generator,
-  Eo as Model,
+  vo as Model,
   fo as NanoGPT,
-  Go as Task,
+  Do as Task,
   lo as TeachableLLM,
-  Ho as checks,
-  Bo as estimateMemoryUsage,
-  bo as estimateParameterCount,
-  yo as estimateResources,
-  Lo as estimateTrainingMemoryUsage,
+  Bo as Trainer,
+  Oo as checks,
+  Lo as estimateMemoryUsage,
+  So as estimateParameterCount,
+  Ao as estimateResources,
+  Fo as estimateTrainingMemoryUsage,
   mo as layers,
   Co as loadTextData,
   eo as ops,
-  No as performanceTest,
-  wo as selectBackend,
-  Jo as sentenceEmbeddings,
-  Ko as sentenceEmbeddingsTensor,
+  Ho as performanceTest,
+  qo as selectBackend,
+  Vo as sentenceEmbeddings,
+  Wo as sentenceEmbeddingsTensor,
   to as tasks,
-  Ro as tokensFromTasks,
-  Ao as topP,
-  So as validateConfig,
+  No as tokensFromTasks,
+  Uo as topP,
+  Go as validateConfig,
   go as waitForModel
 };

package/dist/training/BasicTrainer.js

@@ -52,19 +52,19 @@ class R {
             mixedPrecision: this._mixedPrecision
           },
           d
-        ), s = T(e, l, this.maskedLoss);
+        ), t = T(e, l, this.maskedLoss);
         e.dispose();
-        const g = s.mul(S(this.lossScaling));
-        return s.dispose(), g;
-      }, { value: t, grads: n } = L(o);
+        const g = t.mul(S(this.lossScaling));
+        return t.dispose(), g;
+      }, { value: s, grads: n } = L(o);
       if (a)
         this.model.getProfiler()?.endMemory("Training");
       else {
         this.optimizer.applyGradients(n);
         const e = Object.keys(n);
-        this.model.weightStore.touchVariables(e), this.model.getProfiler()?.endMemory("Training"), c ? (i.gradients = n, Object.values(n).forEach((s) => P(s))) : m(n);
+        this.model.weightStore.touchVariables(e), this.model.getProfiler()?.endMemory("Training"), c ? (i.gradients = n, Object.values(n).forEach((t) => P(t))) : m(n);
       }
-      return t.mul(S(1 / this.lossScaling));
+      return s.mul(S(1 / this.lossScaling));
     });
   }
   async dummyPass() {
@@ -99,13 +99,13 @@ class R {
       ...r
     }, d = Date.now(), l = this.createEmptyState();
     this.lastState = l, await this.dummyPass(), r?.advancedMetrics && (this.model.getProfiler() || this.model.setProfiler(new y())), this.running = !0, l.logStartTime = d;
-    const o = a ? new u(this.model, a) : void 0, t = await i.iterator();
+    const o = a ? new u(this.model, a) : void 0, s = await i.iterator();
     try {
       for (; this.running; ) {
-        const n = await t.next();
+        const n = await s.next();
         if (n.done) break;
-        const e = n.value, s = this.trainStep(l, e, !1);
-        e.xs.dispose(), e.ys.dispose(), l.step++, l.totalSteps++, l.step % c === 0 && await this.performLogging(s, e.xs.shape[0], r, o), s.dispose();
+        const e = n.value, t = this.trainStep(l, e, !1);
+        e.xs.dispose(), e.ys.dispose(), l.step++, l.totalSteps++, l.step % c === 0 && await this.performLogging(t, e.xs.shape[0], r, o), t.dispose();
       }
     } catch (n) {
       throw console.error("Training error:", n), m(), n;
@@ -116,45 +116,45 @@ class R {
     const { onStep: d } = {
       ...p,
       ...a
-    }, l = a?.gradientMetrics || !1, o = (await i.data())[0], t = this.lastState;
-    t.lastLoss = o;
+    }, l = a?.gradientMetrics || !1, o = (await i.data())[0], s = this.lastState;
+    s.lastLoss = o;
     const n = Date.now();
-    t.trainingDuration += n - t.logStartTime;
+    s.trainingDuration += n - s.logStartTime;
     const e = {
-      loss: t.lastLoss,
-      step: t.step,
-      time: Date.now() - t.logStartTime,
+      loss: s.lastLoss,
+      step: s.step,
+      time: Date.now() - s.logStartTime,
       batchSize: r,
       learningRate: a?.advancedMetrics ? this.optimizer.lr : void 0
     };
     if (this.model.trainingState = {
-      steps: t.totalSteps,
+      steps: s.totalSteps,
       learningRate: this.optimizer.lr,
       batchSize: r,
-      loss: t.lastLoss
-    }, a?.gradientMetrics && l && t.gradients) {
-      const s = /* @__PURE__ */ new Map();
-      for (const [g, h] of Object.entries(t.gradients))
-        s.set(g, await k(h)), h.dispose();
-      e.gradientMetrics = s;
+      loss: s.lastLoss
+    }, a?.gradientMetrics && l && s.gradients) {
+      const t = /* @__PURE__ */ new Map();
+      for (const [g, h] of Object.entries(s.gradients))
+        t.set(g, await k(h)), h.dispose();
+      e.gradientMetrics = t;
     }
     if (c)
       try {
-        const s = await c.evaluate(5);
-        t.validationLosses.push(s), e.valLoss = s;
-      } catch (s) {
-        console.error("Validation error:", s);
+        const t = await c.evaluate(5);
+        Array.isArray(t) ? e.valLoss = t[0] : (s.validationLosses.push(t), e.valLoss = t);
+      } catch (t) {
+        console.error("Validation error:", t);
       }
     if (d) {
-      const s = {
-        duration: t.trainingDuration,
-        totalSamples: t.totalSteps * e.batchSize,
-        samplesPerSecond: t.totalSteps * e.batchSize / (t.trainingDuration / 1e3),
+      const t = {
+        duration: s.trainingDuration,
+        totalSamples: s.totalSteps * e.batchSize,
+        samplesPerSecond: s.totalSteps * e.batchSize / (s.trainingDuration / 1e3),
         memory: a?.advancedMetrics ? this.model.getProfiler()?.getPeakMemory() || 0 : void 0
       };
-      await d(e, s);
+      await d(e, t);
     }
-    t.logStartTime = Date.now();
+    s.logStartTime = Date.now();
   }
   async trainOnDataset(i, r, a) {
     const { logInterval: c, maxSteps: d } = {
@@ -162,13 +162,13 @@ class R {
       ...r
     }, l = Date.now(), o = this.createEmptyState();
     this.lastState = o, await this.dummyPass(), r?.advancedMetrics && (this.model.getProfiler() || this.model.setProfiler(new y())), this.running = !0, o.logStartTime = l;
-    const t = a ? new u(this.model, a) : void 0, n = await i.iterator();
+    const s = a ? new u(this.model, a) : void 0, n = await i.iterator();
     try {
       for (; this.running; ) {
         const e = await n.next();
         if (e.done) break;
-        const s = e.value, g = o.step % c === 0, h = (r?.gradientMetrics || !1) && g, f = this.trainStep(o, s, !1, h);
-        s.xs.dispose(), s.ys.dispose(), o.step++, o.totalSteps++, g && await this.performLogging(f, s.xs.shape[0], r, t), f.dispose(), o.step >= d && this.stop();
+        const t = e.value, g = o.step % c === 0, h = (r?.gradientMetrics || !1) && g, f = this.trainStep(o, t, !1, h);
+        t.xs.dispose(), t.ys.dispose(), o.step++, o.totalSteps++, g && await this.performLogging(f, t.xs.shape[0], r, s), f.dispose(), o.step >= d && this.stop();
       }
     } catch (e) {
       throw console.error("Training error:", e), m(), e;

package/dist/training/Evaluator.d.ts

@@ -1,9 +1,14 @@
 import { Dataset } from '@tensorflow/tfjs-data';
 import { TensorContainer } from '@tensorflow/tfjs-core';
 import { default as Model, ModelForwardAttributes } from '../models/model';
+import { Conversation, ITokeniser } from '../main';
 export default class Evaluator {
     private model;
-    private iterator;
-    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer>);
-    evaluate(maxBatches?: number): Promise<number>;
+    private iterator?;
+    private xs?;
+    private ys?;
+    constructor(model: Model<ModelForwardAttributes>, dataset: Dataset<TensorContainer> | Conversation[][], tokeniser?: ITokeniser);
+    dispose(): void;
+    private calculateBatchLoss;
+    evaluate(maxBatches?: number): Promise<number | number[]>;
 }

package/dist/training/Evaluator.js

@@ -1,23 +1,47 @@
-import { calculateLoss as h } from "./loss.js";
-class m {
-  constructor(s, t) {
-    this.model = s, this.iterator = t.iterator();
+import "../index-twYeuV3_.js";
+import { calculateLoss as u } from "./loss.js";
+import { buildSFTExample as p } from "./SFTDatasetBuilder.js";
+import { t as c } from "../tensor-CO6h2H2F.js";
+class b {
+  constructor(i, t, a) {
+    if (this.model = i, Array.isArray(t)) {
+      if (!a)
+        throw new Error("Tokeniser is required when dataset is an array of conversations");
+      const o = t.map((s) => p(s, -100, a, i.config.blockSize)).filter((s) => s !== null);
+      if (o.length === 0)
+        return;
+      this.xs = c(o.map((s) => s.xs)), this.ys = c(o.map((s) => s.ys));
+    } else
+      this.iterator = t.iterator();
   }
   iterator;
-  async evaluate(s = 100) {
-    let t = 0, o = 0;
-    const n = await this.iterator;
-    for (let a = 0; a < s; a++) {
-      const e = await n.next();
-      if (e.done) break;
-      const u = e.value, { xs: r, ys: i } = u, l = this.model.forward({ training: !1 }, r), c = h(l, i);
-      l.dispose(), r.dispose(), i.dispose();
-      const d = await c.array();
-      c.dispose(), t += d, o++;
-    }
-    return t / o;
+  xs;
+  ys;
+  dispose() {
+    this.xs && this.xs.dispose(), this.ys && this.ys.dispose();
+  }
+  async calculateBatchLoss(i, t, a, o) {
+    const s = this.model.forward({ training: !1 }, i), r = u(s, t, o, a);
+    s.dispose();
+    const e = await r.array();
+    return r.dispose(), e;
+  }
+  async evaluate(i = 100) {
+    let t = 0, a = 0;
+    if (this.iterator) {
+      const o = await this.iterator;
+      for (let s = 0; s < i; s++) {
+        const r = await o.next();
+        if (r.done) break;
+        const l = r.value, { xs: e, ys: n } = l, h = this.model.forward({ training: !1 }, e), f = await this.calculateBatchLoss(h, n, !1, !1);
+        e.dispose(), n.dispose(), t += f, a++;
+      }
+      return t / a;
+    } else if (this.xs && this.ys)
+      return this.calculateBatchLoss(this.xs, this.ys, !0, !0);
+    throw new Error("No data available for evaluation");
   }
 }
 export {
-  m as default
+  b as default
 };

package/dist/training/SFTDatasetBuilder.d.ts

@@ -1,7 +1,11 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-import { ITokeniser } from '../tokeniser/type';
+import { Conversation, ITokeniser } from '../tokeniser/type';
 import { Dataset } from '@tensorflow/tfjs-data';
 import { Task } from './tasks/Task';
+export declare function buildSFTExample(conversation: Conversation[], ignoreIndex: number, tokenizer: ITokeniser, blockSize: number): {
+    xs: Int32Array;
+    ys: Int32Array;
+} | null;
 export declare class SFTDatasetBuilder {
     tokenizer: ITokeniser;
     blockSize: number;

package/dist/training/SFTDatasetBuilder.js

@@ -1,41 +1,44 @@
 import { t as x } from "../index-twYeuV3_.js";
 import "../dataset-BMe3pbsL.js";
-import { g as y } from "../readers-C_41Nuv3.js";
+import { g as I } from "../readers-C_41Nuv3.js";
 import "../index-Cp39cXWe.js";
-function I(u, a, t, c) {
-  const s = [t.bosToken], n = [!1], d = {
+function w(u, a, t, l) {
+  const s = [t.bosToken], n = [!1], f = {
     user: t.getSpecialTokenIndex("<|user_start|>"),
     assistant: t.getSpecialTokenIndex("<|assistant_start|>"),
     system: t.getSpecialTokenIndex("<|system_start|>")
-  }, r = {
+  }, i = {
     user: t.getSpecialTokenIndex("<|user_end|>"),
     assistant: t.getSpecialTokenIndex("<|assistant_end|>"),
     system: t.getSpecialTokenIndex("<|system_end|>")
   };
   for (const e of u) {
-    const p = d[e.role], l = r[e.role];
-    if (p == null || l == null)
+    const c = f[e.role], h = i[e.role];
+    if (c == null || h == null)
       throw new Error(`Missing special tokens for role: ${e.role}`);
-    s.push(p), n.push(!1);
-    const T = t.encode(e.content);
-    for (const g of T) {
-      s.push(g);
-      const S = t.isSpecialToken(g), k = e.role === "assistant";
-      n.push(k && !S);
+    s.push(c), n.push(!1);
+    const k = t.encode(e.content);
+    for (const m of k) {
+      s.push(m);
+      const S = t.isSpecialToken(m), y = e.role === "assistant";
+      n.push(y && !S);
     }
-    s.push(l), n.push(!1);
+    s.push(h), n.push(!1);
   }
   s.push(t.eosToken), n.push(!1);
-  const o = c + 1;
+  const o = l + 1;
   if (s.length < o) {
-    const e = o - s.length, p = t.getSpecialTokenIndex("<pad>");
-    for (let l = 0; l < e; l++)
-      s.push(p), n.push(!1);
+    const e = o - s.length, c = t.getSpecialTokenIndex("<pad>");
+    for (let h = 0; h < e; h++)
+      s.push(c), n.push(!1);
   } else s.length > o && (s.length = o, n.length = o);
-  const h = new Int32Array(s.slice(0, c)), i = s.slice(1, c + 1), m = n.slice(1, c + 1), f = new Int32Array(i.length);
-  for (let e = 0; e < i.length; e++)
-    f[e] = m[e] ? i[e] : a;
-  return { xs: h, ys: f };
+  const p = new Int32Array(s.slice(0, l)), r = s.slice(1, l + 1), T = n.slice(1, l + 1), d = new Int32Array(r.length);
+  let g = !1;
+  for (let e = 0; e < r.length; e++) {
+    const c = T[e] ? r[e] : a;
+    d[e] = c, c !== a && (g = !0);
+  }
+  return g ? { xs: p, ys: d } : null;
 }
 class A {
   tokenizer;
@@ -49,17 +52,17 @@ class A {
    * - Pads with eosToken and masks padding.
    * - Masks non-assistant tokens in labels with ignoreIndex (default -100).
    */
-  async createSFTDataset(a, t = 32, c = -100) {
+  async createSFTDataset(a, t = 32, l = -100) {
     if (!a.length)
       throw new Error("No conversations provided.");
     const s = this.tokenizer, n = this.blockSize;
-    return y(function* () {
+    return I(function* () {
       for (; ; ) {
-        const r = Math.floor(Math.random() * a.length), h = a[r].getRandomConversation();
-        yield I(h, c, s, n);
+        const i = Math.floor(Math.random() * a.length), p = a[i].getRandomConversation(), r = w(p, l, s, n);
+        r && (yield r);
       }
-    }).batch(t).map((r) => {
-      const o = r;
+    }).batch(t).map((i) => {
+      const o = i;
       return x(() => ({
         xs: o.xs.cast("int32"),
         ys: o.ys.cast("int32")
@@ -68,5 +71,6 @@ class A {
   }
 }
 export {
-  A as SFTDatasetBuilder
+  A as SFTDatasetBuilder,
+  w as buildSFTExample
 };

package/dist/training/loss.d.ts

@@ -1,2 +1,2 @@
 import { Tensor } from '@tensorflow/tfjs-core';
-export declare function calculateLoss(logits: Tensor, targets: Tensor, masked?: boolean): Tensor;
+export declare function calculateLoss(logits: Tensor, targets: Tensor, masked?: boolean, keepBatch?: boolean): Tensor;

package/dist/training/loss.js

@@ -1,11 +1,11 @@
-import { createSoftmaxCrossEntropyWithGrad as n } from "./sparseCrossEntropy.js";
-function l(r, s, t) {
+import { createSoftmaxCrossEntropyWithGrad as c } from "./sparseCrossEntropy.js";
+function a(r, s, t, n) {
   try {
-    return n(t)(r, s);
+    return c(t, n)(r, s);
   } catch (o) {
     throw console.error("Error computing loss:", o), new Error(`Loss computation failed: ${o}`);
   }
 }
 export {
-  l as calculateLoss
+  a as calculateLoss
 };

package/dist/training/sparseCrossEntropy.d.ts

@@ -3,5 +3,5 @@ import * as tf from '@tensorflow/tfjs-core';
  * Numerically stable sparse cross-entropy with gradient support
  * This version handles potential numerical issues better
  */
-export declare function sparseSoftmaxCrossEntropy(logits: tf.Tensor, labels: tf.Tensor, validMask?: tf.Tensor): tf.Tensor;
-export declare function createSoftmaxCrossEntropyWithGrad(masked?: boolean): (...args: tf.Tensor[]) => tf.Tensor<tf.Rank>;
+export declare function sparseSoftmaxCrossEntropy(logits: tf.Tensor, labels: tf.Tensor, validMask?: tf.Tensor, keepBatch?: boolean, originalBatchShape?: number[]): tf.Tensor;
+export declare function createSoftmaxCrossEntropyWithGrad(masked?: boolean, keepBatch?: boolean): (...args: tf.Tensor[]) => tf.Tensor<tf.Rank>;

package/dist/training/sparseCrossEntropy.js

@@ -1,42 +1,46 @@
-import { gatherSub as T } from "../ops/gatherSub.js";
-import { scatterSub as j } from "../ops/scatterSub.js";
-import { t as L, c as k, m as C, j as q, a1 as B, a as y, a2 as v } from "../index-twYeuV3_.js";
-import { m as F, l as P, a as V, n as W, w as $ } from "../not_equal-DXJHGhGS.js";
-import { s as b } from "../sum-CgGUPVhu.js";
-import { s as A } from "../softmax-DpG1TdjZ.js";
-function D(a, p, d) {
-  return L(() => {
-    const o = a.shape[a.shape.length - 1], h = a.shape.slice(0, -1).reduce((t, S) => t * S, 1), i = a.shape.length > 2 ? a.reshape([h, o]) : a, g = p.shape.length > 1 ? p.reshape([h]).cast("int32") : p.cast("int32"), l = F(i, -1, !0), r = k(i, l), n = P(r, -1);
-    let s = T(n, g, r);
-    if (d) {
-      s = C(s, d);
-      const t = b(d);
-      s = q(b(s), t);
-    } else
-      s = V(s);
+import { gatherSub as P } from "../ops/gatherSub.js";
+import { scatterSub as V } from "../ops/scatterSub.js";
+import { t as z, c as W, m as w, j as C, a1 as $, a as L, a2 as E } from "../index-twYeuV3_.js";
+import { m as A, l as D, a as G, n as H, w as J } from "../not_equal-DXJHGhGS.js";
+import { s as m } from "../sum-CgGUPVhu.js";
+import { s as K } from "../softmax-DpG1TdjZ.js";
+function N(e, d, h, x, a) {
+  return z(() => {
+    const u = e.shape[e.shape.length - 1], c = a || e.shape.slice(0, -1), f = c.reduce((n, y) => n * y, 1), i = e.shape.length > 2 ? e.reshape([f, u]) : e, S = d.shape.length > 1 ? d.reshape([f]).cast("int32") : d.cast("int32"), p = A(i, -1, !0), t = W(i, p), r = D(t, -1);
+    let s = P(r, S, t);
+    if (h)
+      if (s = w(s, h), x) {
+        const n = m(h.reshape(c), -1);
+        s = C(m(s.reshape(c), -1), n);
+      } else {
+        const n = m(h);
+        s = C(m(s), n);
+      }
+    else
+      x ? s = G(s.reshape(c), -1) : s = G(s);
     return s;
   });
 }
-function Y(a) {
-  return B(
+function k(e, d) {
+  return $(
     // @ts-expect-error Invalid params
-    (o, u, h) => {
-      const i = o.shape[o.shape.length - 1], l = o.shape.slice(0, -1).reduce((c, e) => c * e, 1), r = o.reshape([l, i]), n = u.reshape([l]).cast("int32");
-      let s, t = null;
-      if (a) {
-        const c = y(-100, "int32"), e = W(n, c);
-        t = e.cast("float32"), s = $(e, n, v(n));
+    (a, u, c) => {
+      const f = a.shape[a.shape.length - 1], i = a.shape.slice(0, -1), S = i.reduce((l, o) => l * o, 1), p = a.reshape([S, f]), t = u.reshape([S]).cast("int32");
+      let r, s = null;
+      if (e) {
+        const l = L(-100, "int32"), o = H(t, l);
+        s = o.cast("float32"), r = J(o, t, E(t));
       } else
-        s = n;
-      const S = D(r, s, t || void 0);
-      return h(t ? [r, s, t] : [r, s]), r.dispose(), n.dispose(), { value: S, gradFunc: (c, e) => L(() => {
-        const f = e[0], E = e[1], m = a ? e[2] : void 0, G = A(f), z = m ? b(m) : y(f.shape[0], "float32"), x = c.div(z).broadcastTo([f.shape[0]]), w = m && a ? C(x, m) : x, I = j(G, E, w), M = v(u);
-        return [I.reshape(o.shape), M];
+        r = t;
+      const n = N(p, r, s || void 0, d, i);
+      return c(s ? [p, r, s] : [p, r]), p.dispose(), t.dispose(), { value: n, gradFunc: (l, o) => z(() => {
+        const b = o[0], I = o[1], g = e ? o[2] : void 0, T = K(b), j = g ? m(g) : L(b.shape[0], "float32"), v = l.div(j).broadcastTo([b.shape[0]]), q = g && e ? w(v, g) : v, F = V(T, I, q), M = E(u);
+        return [F.reshape(a.shape), M];
       }) };
     }
   );
 }
 export {
-  Y as createSoftmaxCrossEntropyWithGrad,
-  D as sparseSoftmaxCrossEntropy
+  k as createSoftmaxCrossEntropyWithGrad,
+  N as sparseSoftmaxCrossEntropy
 };

package/package.json

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