@genai-fi/nanogpt 0.2.12 → 0.3.0

package/dist/training/Evaluator.d.ts

@@ -1,8 +1,9 @@
 import { default as NanoGPT } from '../NanoGPTModel';
-import { default as TF } from '@tensorflow/tfjs';
+import { Dataset } from '@tensorflow/tfjs-data';
+import { TensorContainer } from '@tensorflow/tfjs-core';
 export default class Evaluator {
     private model;
     private iterator;
-    constructor(model: NanoGPT, dataset: TF.data.Dataset<TF.TensorContainer>);
+    constructor(model: NanoGPT, dataset: Dataset<TensorContainer>);
     evaluate(maxBatches?: number): Promise<number>;
 }

package/dist/training/FullTrainer.d.ts

@@ -1,15 +1,16 @@
 import { ITokeniser } from '../tokeniser/type';
 import { default as NanoGPT } from '../NanoGPTModel';
-import { default as TF } from '@tensorflow/tfjs';
 import { default as GPTTrainer, TrainingOptions } from './Trainer';
+import { Tensor } from '@tensorflow/tfjs-core';
+import { Dataset } from '@tensorflow/tfjs-data';
 export default class FullTrainer extends GPTTrainer {
-    constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
-    trainOnDataset(dataset: TF.data.Dataset<{
-        xs: TF.Tensor;
-        ys: TF.Tensor;
-    }>, options: Partial<TrainingOptions>, validationDataset?: TF.data.Dataset<{
-        xs: TF.Tensor;
-        ys: TF.Tensor;
+    constructor(model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    trainOnDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
     }>): Promise<{
         losses: number[];
         validationLosses: number[];

package/dist/training/FullTrainer.js

@@ -1,19 +1,20 @@
-import { generateText as L } from "../utilities/generate.js";
-import w from "./Trainer.js";
-import x from "./Evaluator.js";
-const g = {
+import { generateText as w } from "../utilities/generate.js";
+import x from "./Trainer.js";
+import g from "./Evaluator.js";
+import { a as h } from "../index-pWA4_lUh.js";
+const S = {
   desiredLoss: 0.01,
   logInterval: 1,
   maxSteps: 1e3
 };
-class P extends w {
-  constructor(r, i, o, n = 3e-4) {
-    super(r, i, o, n);
+class D extends x {
+  constructor(r, i, o = 3e-4) {
+    super(r, i, o);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
   async trainOnDataset(r, i, o) {
-    const { desiredLoss: n, logInterval: m, onStep: l, prompt: c, maxSteps: d } = {
-      ...g,
+    const { desiredLoss: m, logInterval: u, onStep: n, prompt: l, maxSteps: d } = {
+      ...S,
       ...i
     }, t = {
       step: 0,
@@ -24,45 +25,45 @@ class P extends w {
       ...this.lastState || {}
     };
     this.lastState = t, this.dummyPass(), this.model.trainable = !0;
-    const u = Date.now();
+    const f = Date.now();
     this.running = !0;
-    const h = o ? new x(this.model, o) : void 0, f = await r.iterator();
+    const c = o ? new g(this.model, o) : void 0, v = await r.iterator();
     try {
-      for (; this.running && !(t.lastLoss < n); ) {
-        const e = await f.next();
+      for (; this.running && !(t.lastLoss < m); ) {
+        const e = await v.next();
         if (e.done) break;
-        const p = e.value, v = this.trainBatch(t, p), a = {
+        const p = e.value, L = this.trainBatch(t, p), a = {
           loss: t.lastLoss,
           step: t.step,
-          time: Date.now() - u,
+          time: Date.now() - f,
           batchSize: p.xs.shape[0]
         };
-        if (this.model.log.push(a), t.step % m === 0) {
-          if (await v, h)
+        if (this.model.log.push(a), t.step % u === 0) {
+          if (await L, c)
             try {
-              const s = await h.evaluate(5);
+              const s = await c.evaluate(5);
               t.validationLosses.push(s), a.valLoss = s;
             } catch (s) {
               console.error("Validation error:", s);
             }
-          if (l) {
-            if (c) {
-              const s = await L(this.tokenizer, this.model, c, 100, {
+          if (n) {
+            if (l) {
+              const s = await w(this.tokenizer, this.model, l, 100, {
                 temperature: 0.8
               });
               a.example = s;
             }
-            await l(a);
+            await n(a);
           }
         }
         t.step >= d && this.stop();
       }
     } catch (e) {
-      throw console.error("Training error:", e), this.tf.dispose(), e;
+      throw console.error("Training error:", e), h(), e;
     }
-    return this.tf.dispose(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
+    return h(), this.running = !1, { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  P as default
+  D as default
 };

package/dist/training/LayerTrainer.d.ts

@@ -1,7 +1,8 @@
 import { ITokeniser } from '../tokeniser/type';
 import { default as NanoGPT } from '../NanoGPTModel';
-import { default as TF } from '@tensorflow/tfjs';
 import { default as GPTTrainer, TrainingOptions } from './Trainer';
+import { Tensor } from '@tensorflow/tfjs-core';
+import { Dataset } from '@tensorflow/tfjs-data';
 interface LayerTrainingOptions extends TrainingOptions {
     stepsPerLayer: number;
     maxPasses: number;
@@ -12,14 +13,14 @@ export default class LayerTrainer extends GPTTrainer {
     private trainingPattern;
     private startPass;
     private startLayer;
-    constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    constructor(model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
     private applyTrainingPattern;
-    trainOnDataset(dataset: TF.data.Dataset<{
-        xs: TF.Tensor;
-        ys: TF.Tensor;
-    }>, options: Partial<LayerTrainingOptions>, validationDataset?: TF.data.Dataset<{
-        xs: TF.Tensor;
-        ys: TF.Tensor;
+    trainOnDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<LayerTrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
     }>): Promise<{
         losses: number[];
         validationLosses: number[];

package/dist/training/LayerTrainer.js

@@ -1,33 +1,34 @@
-import { generateText as S } from "../utilities/generate.js";
-import u from "./Trainer.js";
-import { schedule as v } from "./lwSchedule.js";
-import w from "./Evaluator.js";
-const T = {
+import { generateText as u } from "../utilities/generate.js";
+import v from "./Trainer.js";
+import { schedule as w } from "./lwSchedule.js";
+import T from "./Evaluator.js";
+import { a as m } from "../index-pWA4_lUh.js";
+const x = {
   desiredLoss: 0.01,
   logInterval: 1,
   stepsPerLayer: 400,
   maxPasses: 3,
   maxSteps: 1e3
 };
-class z extends u {
+class E extends v {
   trainingPattern = [];
   startPass = 0;
   startLayer = 0;
-  constructor(r, s, e, p = 3e-4) {
-    if (super(r, s, e, p), this.trainingPattern = v[s.config.nLayer - 1] || [], s.log.length > 0) {
+  constructor(s, r, e = 3e-4) {
+    if (super(s, r, e), this.trainingPattern = w[s.config.nLayer - 1] || [], s.log.length > 0) {
       const i = s.log[s.log.length - 1];
       i.pass !== void 0 && i.layer !== void 0 && (this.startPass = i.pass, this.startLayer = i.layer, console.log(`Resuming training from pass ${this.startPass}, layer ${this.startLayer}`));
     }
   }
-  applyTrainingPattern(r) {
-    const s = r < this.trainingPattern.length ? r : this.trainingPattern.length - 1, e = this.trainingPattern[s];
-    this.model.setSkipMask(e.skip), this.model.setTrainableMask(e.trainable), this.resetOptimizer(e.adam), console.log("Applied training pattern:", s, e);
+  applyTrainingPattern(s) {
+    const r = s < this.trainingPattern.length ? s : this.trainingPattern.length - 1, e = this.trainingPattern[r];
+    this.model.setSkipMask(e.skip), this.model.setTrainableMask(e.trainable), this.resetOptimizer(e.adam), console.log("Applied training pattern:", r, e);
   }
   // Train for multiple epochs using Dataset API - FIXED memory leaks
-  async trainOnDataset(r, s, e) {
-    const { desiredLoss: p, logInterval: i, stepsPerLayer: L, onLayerChange: o, onPassComplete: h, onStep: c, prompt: g } = {
-      ...T,
-      ...s
+  async trainOnDataset(s, r, e) {
+    const { desiredLoss: i, logInterval: L, stepsPerLayer: f, onLayerChange: o, onPassComplete: p, onStep: h, prompt: c } = {
+      ...x,
+      ...r
     }, t = {
       pass: 0,
       layerStep: 0,
@@ -39,51 +40,51 @@ class z extends u {
       validationLosses: []
     };
     this.dummyPass();
-    const f = Date.now();
+    const d = Date.now();
     this.startPass = 0, this.startLayer = 0;
-    const y = e ? new w(this.model, e) : void 0, d = await r.iterator();
+    const g = e ? new T(this.model, e) : void 0, P = await s.iterator();
     this.applyTrainingPattern(t.layerStep % this.trainingPattern.length);
     try {
-      for (; !(t.lastLoss < p); ) {
-        const n = await d.next();
+      for (; !(t.lastLoss < i); ) {
+        const n = await P.next();
         if (n.done) break;
-        const m = n.value, P = this.trainBatch(t, m);
+        const y = n.value, S = this.trainBatch(t, y);
         t.stepSinceLayerChange++;
         const l = {
           loss: t.lastLoss,
           step: t.step,
-          time: Date.now() - f,
-          batchSize: m.xs.shape[0],
+          time: Date.now() - d,
+          batchSize: y.xs.shape[0],
           pass: t.pass,
           layer: t.layerStep % this.model.config.nLayer
         };
-        if (this.model.log.push(l), t.step % i === 0) {
-          if (await P, y)
+        if (this.model.log.push(l), t.step % L === 0) {
+          if (await S, g)
             try {
-              const a = await y.evaluate(5);
+              const a = await g.evaluate(5);
               t.validationLosses.push(a), l.valLoss = a;
             } catch (a) {
               console.error("Validation error:", a);
             }
-          if (c) {
-            if (g) {
-              const a = await S(this.tokenizer, this.model, g, 100, {
+          if (h) {
+            if (c) {
+              const a = await u(this.tokenizer, this.model, c, 100, {
                 temperature: 0.8,
                 topK: 10
               });
               l.example = a;
             }
-            await c(l);
+            await h(l);
           }
         }
-        t.stepSinceLayerChange >= L && (t.layerStep++, t.layerStep % this.model.config.nLayer === 0 ? (o && await o(t.layerStep, t.pass), h && await h(t.pass), t.pass++) : o && await o(t.layerStep, t.pass), t.stepSinceLayerChange = 0, this.applyTrainingPattern(t.layerStep % this.trainingPattern.length));
+        t.stepSinceLayerChange >= f && (t.layerStep++, t.layerStep % this.model.config.nLayer === 0 ? (o && await o(t.layerStep, t.pass), p && await p(t.pass), t.pass++) : o && await o(t.layerStep, t.pass), t.stepSinceLayerChange = 0, this.applyTrainingPattern(t.layerStep % this.trainingPattern.length));
       }
     } catch (n) {
-      throw console.error("Training error:", n), this.tf.dispose(), n;
+      throw console.error("Training error:", n), m(), n;
     }
-    return this.tf.dispose(), { losses: t.losses, validationLosses: t.validationLosses };
+    return m(), { losses: t.losses, validationLosses: t.validationLosses };
   }
 }
 export {
-  z as default
+  E as default
 };

package/dist/training/Trainer.d.ts

@@ -1,8 +1,10 @@
 import { ITokeniser } from '../tokeniser/type';
 import { DatasetBuilder } from './DatasetBuilder';
 import { default as NanoGPT, TrainingLogEntry } from '../NanoGPTModel';
-import { default as TF } from '@tensorflow/tfjs';
 import { default as AdamExt } from './AdamExt';
+import { TensorContainer } from '@tensorflow/tfjs-core/dist/tensor_types';
+import { Scalar, Tensor } from '@tensorflow/tfjs-core';
+import { Dataset } from '@tensorflow/tfjs-data';
 export interface TrainingState {
     step: number;
     lastLoss: number;
@@ -28,11 +30,10 @@ export default abstract class GPTTrainer {
     protected model: NanoGPT;
     protected optimizer: AdamExt;
     protected datasetBuilder: DatasetBuilder;
-    protected tf: typeof TF;
     protected learningRate: number;
     protected running: boolean;
     protected lastState?: TrainingState;
-    constructor(tf: typeof TF, model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
+    constructor(model: NanoGPT, tokenizer: ITokeniser, learningRate?: number);
     setLearningRate(learningRate: number): void;
     reset(): void;
     stop(): void;
@@ -40,34 +41,34 @@ export default abstract class GPTTrainer {
     resetOptimizer(config?: AdamConfig): void;
     private printGradients;
     protected trainStep(batch: {
-        xs: TF.Tensor;
-        ys: TF.Tensor;
-    }, dummy?: boolean, print?: boolean): TF.Scalar;
+        xs: Tensor;
+        ys: Tensor;
+    }, dummy?: boolean, print?: boolean): Scalar;
     protected dummyPass(): void;
     protected trainBatch(state: TrainingState, batch: {
-        xs: TF.Tensor;
-        ys: TF.Tensor;
+        xs: Tensor;
+        ys: Tensor;
     }): Promise<number>;
-    abstract trainOnDataset(dataset: TF.data.Dataset<{
-        xs: TF.Tensor;
-        ys: TF.Tensor;
-    }>, options: Partial<TrainingOptions>, validationDataset?: TF.data.Dataset<{
-        xs: TF.Tensor;
-        ys: TF.Tensor;
+    abstract trainOnDataset(dataset: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
+    }>, options: Partial<TrainingOptions>, validationDataset?: Dataset<{
+        xs: Tensor;
+        ys: Tensor;
     }>): Promise<{
         losses: number[];
         validationLosses: number[];
     }>;
     createTrainValidationSplit(textData: string[], batchSize?: number, validationSplit?: number): Promise<{
-        trainDataset: TF.data.Dataset<{
-            xs: TF.Tensor;
-            ys: TF.Tensor;
+        trainDataset: Dataset<{
+            xs: Tensor;
+            ys: Tensor;
         }>;
-        validationDataset: TF.data.Dataset<{
-            xs: TF.Tensor;
-            ys: TF.Tensor;
+        validationDataset: Dataset<{
+            xs: Tensor;
+            ys: Tensor;
         }>;
     }>;
-    createDataset(textData: string[], batchSize?: number): Promise<TF.data.Dataset<TF.TensorContainer>>;
+    createDataset(textData: string[], batchSize?: number): Promise<Dataset<TensorContainer>>;
     dispose(): void;
 }

package/dist/training/Trainer.js

@@ -1,13 +1,16 @@
-import { DatasetBuilder as d } from "./DatasetBuilder.js";
+import { DatasetBuilder as p } from "./DatasetBuilder.js";
 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);
+import { t as g, v as u, a as o } from "../index-pWA4_lUh.js";
+import { m as y, a as f, b as S, n as z } from "../moments-DfcpfwKi.js";
+import { m as b } from "../sum-C7Mgy9Bw.js";
+import { z as n } from "../zeros-CCy9C3uU.js";
+class M {
+  constructor(t, s, e = 1e-3) {
+    this.tokenizer = s, this.model = t, this.learningRate = e, this.resetOptimizer(), this.datasetBuilder = new p(s, t.config.blockSize);
   }
   model;
   optimizer;
   datasetBuilder;
-  tf;
   learningRate;
   running = !1;
   lastState;
@@ -42,21 +45,21 @@ class g {
   printGradients(t) {
     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]}`);
+      console.log(`${s}:`), console.log(`  Shape: ${e.shape}`), console.log(`  Mean: ${y(e).dataSync()[0]}`), console.log(`  Std: ${f(e).variance.sqrt().dataSync()[0]}`), console.log(`  Min: ${S(e).dataSync()[0]}`), console.log(`  Max: ${b(e).dataSync()[0]}`), console.log(`  Norm: ${z(e).dataSync()[0]}`);
     });
   }
   trainStep(t, s = !1, e = !1) {
-    return this.tf.tidy(() => {
+    return g(() => {
       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 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;
+      const { xs: a, ys: r } = t, l = () => {
+        const { loss: d, logits: m } = this.model.forward(a, r, !0);
+        return m.dispose(), d;
+      }, { value: c, grads: i } = u(l);
+      return s ? this.model.getProfiler()?.endMemory("Training") : (e && (console.log("-------"), this.printGradients(i), console.log("-------")), this.optimizer.applyGradients(i), this.model.getProfiler()?.endMemory("Training"), o(i)), c;
     });
   }
   dummyPass() {
-    const t = this.tf.zeros([1, this.model.config.blockSize], "int32"), s = this.tf.zeros([1, this.model.config.blockSize], "int32");
+    const t = n([1, this.model.config.blockSize], "int32"), s = n([1, this.model.config.blockSize], "int32");
     try {
       const e = this.trainStep({ xs: t, ys: s }, !0);
       e.dataSync(), e.dispose();
@@ -69,19 +72,19 @@ class g {
   async trainBatch(t, s) {
     try {
       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));
+      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), this.tf.dispose(), e;
+      throw console.error(`Error processing batch at step ${t.step}:`, e), o(), e;
     }
   }
   async createTrainValidationSplit(t, s = 32, e = 0.1) {
-    const i = await this.datasetBuilder.createTextDataset(t, s, 0, 1 - e), a = await this.datasetBuilder.createTextDataset(
+    const a = await this.datasetBuilder.createTextDataset(t, s, 0, 1 - e), r = await this.datasetBuilder.createTextDataset(
       t,
       s,
       1 - e,
       1
     );
-    return { trainDataset: i, validationDataset: a };
+    return { trainDataset: a, validationDataset: r };
   }
   async createDataset(t, s = 32) {
     return await this.datasetBuilder.createTextDataset(t, s);
@@ -91,5 +94,5 @@ class g {
   }
 }
 export {
-  g as default
+  M as default
 };

package/dist/training/sparseCrossEntropy.js

@@ -1,178 +1,34 @@
-import { gatherSub as _ } from "../ops/gatherSub.js";
-import { scatterSub as v } from "../ops/scatterSub.js";
-import { o as l, d, E as f, M as w, y as I, L as z, z as N, b as E, D as F, F as M, e as m, H as g, I as $, J as S } from "../index-YPKosni4.js";
-import { r as b } from "../reshape-DmnmKT6r.js";
-import { s as T } from "../sum-D7fu15XL.js";
-/**
- * @license
- * Copyright 2017 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.
- * =============================================================================
- */
-function D(n, s, t) {
-  const a = n.length + s.length, e = [];
-  let r = 0, c = 0;
-  for (let o = 0; o < a; o++)
-    t.indexOf(o) === -1 ? e.push(n[r++]) : e.push(s[c++]);
-  return e;
-}
-function P(n, s) {
-  const t = s.map((a) => 1);
-  return D(n, t, s);
-}
-/**
- * @license
- * Copyright 2020 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.
- * =============================================================================
- */
-function A(n, s = null, t = !1) {
-  const e = { x: d(n, "x", "max") }, r = { reductionIndices: s, keepDims: t };
-  return f.runKernel(w, e, r);
-}
-const L = /* @__PURE__ */ l({ max_: A });
-/**
- * @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.
- * =============================================================================
- */
-function H(n) {
-  const t = { x: d(n, "x", "exp") };
-  return f.runKernel(I, t);
-}
-const J = /* @__PURE__ */ l({ exp_: H });
-/**
- * @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.
- * =============================================================================
- */
-function O(n) {
-  const t = { x: d(n, "x", "log", "float32") };
-  return f.runKernel(z, t);
-}
-const W = /* @__PURE__ */ l({ log_: O });
-/**
- * @license
- * Copyright 2020 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.
- * =============================================================================
- */
-function j(n, s = null, t = !1) {
-  const a = d(n, "x", "logSumExp"), e = N(s, a.shape), r = L(
-    a,
-    e,
-    !0
-    /* keepDims */
-  ), c = E(a, r), o = J(c), p = T(o, e), u = W(p), x = F(b(r, u.shape), u);
-  if (t) {
-    const h = P(x.shape, e);
-    return b(x, h);
-  }
-  return x;
-}
-const q = /* @__PURE__ */ l({ logSumExp_: j });
-/**
- * @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.
- * =============================================================================
- */
-function B(n, s = -1) {
-  const t = d(n, "logits", "softmax", "float32");
-  if (s === -1 && (s = t.rank - 1), s !== t.rank - 1)
-    throw Error(`Softmax along a non-last dimension is not yet supported. Logits was rank ${t.rank} and dim was ${s}`);
-  const a = { logits: t }, e = { dim: s };
-  return f.runKernel(M, a, e);
-}
-const Q = /* @__PURE__ */ l({ softmax_: B });
-function R(n, s) {
-  return $(() => {
-    const t = n.shape[n.shape.length - 1], e = n.shape.slice(0, -1).reduce((h, i) => h * i, 1), r = n.shape.length > 2 ? n.reshape([e, t]) : n, c = s.shape.length > 1 ? s.reshape([e]).cast("int32") : s.cast("int32"), o = L(r, -1, !0), p = E(r, o), u = q(p, -1);
-    return _(u, c, p);
+import { gatherSub as L } from "../ops/gatherSub.js";
+import { scatterSub as y } from "../ops/scatterSub.js";
+import { e as u, c as i, z as S, t as f, s as G } from "../index-pWA4_lUh.js";
+import { s as v, l as z } from "../softmax-Be_lsqUc.js";
+import { m as k } from "../sum-C7Mgy9Bw.js";
+function F(a, s) {
+  return f(() => {
+    const e = a.shape[a.shape.length - 1], n = a.shape.slice(0, -1).reduce((d, c) => d * c, 1), p = a.shape.length > 2 ? a.reshape([n, e]) : a, o = s.shape.length > 1 ? s.reshape([n]).cast("int32") : s.cast("int32"), t = k(p, -1, !0), r = G(p, t), h = z(r, -1);
+    return L(h, o, r);
   });
 }
-function ts() {
-  return m().backendName === "tensorflow" ? g((s, t, a) => {
-    const e = s.shape.length > 2 ? s.reshape([-1, s.shape[s.shape.length - 1]]) : s, r = t.shape.length > 1 ? t.reshape([-1]).cast("int32") : t.cast("int32"), [c, o] = m().runKernel(
+function $() {
+  return u().backendName === "tensorflow" ? i((s, e, m) => {
+    const n = 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"), [o, t] = u().runKernel(
       "NativeSparseSoftmaxCrossEntropy",
-      { logits: e, labels: r },
+      { logits: n, labels: p },
       {}
     );
-    return a([o.reshape(s.shape)]), { value: c, gradFunc: (p, u) => [u[0], S(t)] };
-  }) : g(
+    return m([t.reshape(s.shape)]), { value: o, gradFunc: (r, h) => [h[0], S(e)] };
+  }) : i(
     // @ts-expect-error Invalid params
-    (s, t, a) => {
-      const e = s.shape[s.shape.length - 1], c = s.shape.slice(0, -1).reduce((h, i) => h * i, 1), o = s.reshape([c, e]), p = t.reshape([c]).cast("int32"), u = R(o, p);
-      return a([o, p]), o.dispose(), p.dispose(), { value: u, gradFunc: (h, i) => $(() => {
-        const y = i[0], k = i[1], C = Q(y), G = v(C, k, h), K = S(t);
-        return [G.reshape(s.shape), K];
+    (s, e, m) => {
+      const n = s.shape[s.shape.length - 1], o = s.shape.slice(0, -1).reduce((d, c) => d * c, 1), t = s.reshape([o, n]), r = e.reshape([o]).cast("int32"), h = F(t, r);
+      return m([t, r]), t.dispose(), r.dispose(), { value: h, gradFunc: (d, c) => f(() => {
+        const l = c[0], b = c[1], x = v(l), C = y(x, b, d), E = S(e);
+        return [C.reshape(s.shape), E];
       }) };
     }
   );
 }
 export {
-  ts as createSoftmaxCrossEntropyWithGrad,
-  R as sparseSoftmaxCrossEntropy
+  $ as createSoftmaxCrossEntropyWithGrad,
+  F as sparseSoftmaxCrossEntropy
 };