@genai-fi/nanogpt 0.9.0 → 0.10.0

package/dist/{RealDiv-D4EzDsC0.js → RealDiv-DgA3z9oO.js}

@@ -1,26 +1,12 @@
-import { aq as T, ag as E, p as O, j as V, aB as B, a1 as F, ah as j, aC as K } from "./index-BzFyqcy-.js";
-import { r as $ } from "./Reshape-Bowtk9BP.js";
-import { g as A, a as C, b as k, c as N, e as R } from "./axis_util-TbGYJ208.js";
-import { t as U, m as W } from "./shared-DuP7ue-R.js";
-import { c as _ } from "./backend_util-CJIiDoV1.js";
-import { f as y } from "./gpgpu_math-CDaYiyE_.js";
-import { g as G, b as L } from "./kernel_funcs_utils-DKLK0Mg3.js";
-/**
- * @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.
- * =============================================================================
- */
+import "./index-ZyQhjEPo.js";
+import { r as $ } from "./Reshape-_kILl6tK.js";
+import { _ as T, g as E, y as B, $ as F } from "./tensor_util-DV-FP5Q3.js";
+import { H as _, e as K, p as O, s as V } from "./tensor-DdQUJZlz.js";
+import { a as A, b as k, d as C, c as N, e as R } from "./axis_util-BvHEw88j.js";
+import { t as U, m as W } from "./shared-BRksrJb3.js";
+import { c as j } from "./backend_util-D-rUb2ty.js";
+import { f as y } from "./gpgpu_math-DDVJCn6-.js";
+import { g as G, b as L } from "./kernel_funcs_utils-Dg_-E44D.js";
 class w {
   constructor(s, e) {
     this.variableNames = ["x"];
@@ -30,7 +16,7 @@ class w {
     let o = "sumValue += dot(values, ones);";
     if (e != null) {
       const p = 1 / e;
-      o = `sumValue += dot(values * ${T(p) ? p.toPrecision(2) : p}, ones);`;
+      o = `sumValue += dot(values * ${_(p) ? p.toPrecision(2) : p}, ones);`;
     }
     let u = "";
     l % t > 0 && (u = `
@@ -89,23 +75,7 @@ class w {
     `;
   }
 }
-/**
- * @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.
- * =============================================================================
- */
-class q {
+class X {
   constructor(s, e) {
     this.variableNames = ["x"];
     const { windowSize: t, batchSize: n, inSize: l, outSize: r } = s;
@@ -213,26 +183,10 @@ class q {
     `;
   }
 }
-/**
- * @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 X(a) {
+function q(a) {
   const s = [];
   for (; s.length === 0 || s[s.length - 1].outSize !== 1; ) {
-    const e = s.length ? s[s.length - 1].outSize : a[1], t = _(e);
+    const e = s.length ? s[s.length - 1].outSize : a[1], t = j(e);
     s.push({
       inSize: e,
       windowSize: t,
@@ -242,39 +196,23 @@ function X(a) {
   return s;
 }
 function P(a, s, e, t) {
-  const n = X(a.shape);
+  const n = q(a.shape);
   let l = a;
   for (let r = 0; r < n.length; r++) {
     const { inSize: i, windowSize: c, outSize: o } = n[r];
     let u, p;
-    e === "mean" ? u = r === 0 ? new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, i) : new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }) : u = new q({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, e), p = l, l = t.runWebGLProgram(u, [l], s), p.dataId !== a.dataId && t.disposeIntermediateTensorInfo(p);
+    e === "mean" ? u = r === 0 ? new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, i) : new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }) : u = new X({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, e), p = l, l = t.runWebGLProgram(u, [l], s), p.dataId !== a.dataId && t.disposeIntermediateTensorInfo(p);
   }
   return l;
 }
-/**
- * @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.
- * =============================================================================
- */
-class Y {
+class H {
   constructor(s, e) {
     this.variableNames = ["A"];
     const t = new Array(s.length);
     for (let r = 0; r < t.length; r++)
       t[r] = s[e[r]];
     this.outputShape = t, this.rank = t.length;
-    const n = y(this.rank), l = H(e);
+    const n = y(this.rank), l = Y(e);
     this.userCode = `
     void main() {
       ${n} resRC = getOutputCoords();
@@ -283,7 +221,7 @@ class Y {
     `;
   }
 }
-function H(a) {
+function Y(a) {
   const s = a.length;
   if (s > 6)
     throw Error(`Transpose for rank ${s} is not yet supported`);
@@ -292,22 +230,6 @@ function H(a) {
     t[a[n]] = e[n];
   return t.join();
 }
-/**
- * @license
- * Copyright 2019 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.
- * =============================================================================
- */
 class J {
   constructor(s, e) {
     this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0;
@@ -340,115 +262,35 @@ class J {
     `;
   }
 }
-/**
- * @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 D(a, s, e) {
-  const t = E().getBool("WEBGL_PACK_ARRAY_OPERATIONS") ? new J(a.shape, s) : new Y(a.shape, s);
+  const t = K().getBool("WEBGL_PACK_ARRAY_OPERATIONS") ? new J(a.shape, s) : new H(a.shape, s);
   return e.runWebGLProgram(t, [a], a.dtype);
 }
-/**
- * @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 Q(a, s, e, t) {
   const n = s, l = a.shape.length, r = O(n, a.shape);
   let i = r;
   const c = A(i, l), o = c != null;
   let u = a;
-  o && (u = D(a, c, t), i = C(i.length, l)), k("sum", i, l);
+  o && (u = D(a, c, t), i = k(i.length, l)), C("sum", i, l);
   const [p, h] = N(u.shape, i);
   let d = p;
   e && (d = R(p, r));
-  const f = V(h), g = V(a.shape) / f, x = $({ inputs: { x: u }, attrs: { shape: [g, f] }, backend: t }), b = B(a.dtype), I = P(x, b, "sum", t), m = $({ inputs: { x: I }, attrs: { shape: d }, backend: t });
+  const f = V(h), g = V(a.shape) / f, x = $({ inputs: { x: u }, attrs: { shape: [g, f] }, backend: t }), b = T(a.dtype), I = P(x, b, "sum", t), m = $({ inputs: { x: I }, attrs: { shape: d }, backend: t });
   return t.disposeIntermediateTensorInfo(x), t.disposeIntermediateTensorInfo(I), o && t.disposeIntermediateTensorInfo(u), m;
 }
-/**
- * @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 Z(a) {
   const { inputs: s, backend: e, attrs: t } = a, { x: n } = s, { axis: l, keepDims: r } = t;
   return Q(n, l, r, e);
 }
-const pe = {
-  kernelName: F,
+const fe = {
+  kernelName: E,
   backendName: "webgl",
   kernelFunc: Z
 };
-/**
- * @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 ee(a, s, e, t) {
   const n = V(s), r = V(a.shape) / n, i = $({ inputs: { x: a }, attrs: { shape: [r, n] }, backend: t }), c = P(i, a.dtype, "max", t), o = $({ inputs: { x: c }, attrs: { shape: e }, backend: t });
   return t.disposeIntermediateTensorInfo(i), t.disposeIntermediateTensorInfo(c), 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 te(a) {
   const { inputs: s, backend: e, attrs: t } = a, { x: n } = s, { reductionIndices: l, keepDims: r } = t, i = n.shape.length, c = O(l, n.shape);
   let o = c;
@@ -465,9 +307,9 @@ function te(a) {
       M.values = z;
     } else
       d = D(n, u, e);
-    o = C(o.length, i);
+    o = k(o.length, i);
   }
-  k("max", o, i);
+  C("max", o, i);
   const [f, S] = N(d.shape, o);
   let g = f;
   r && (g = R(f, c));
@@ -481,27 +323,11 @@ function te(a) {
     x = ee(d, S, g, e);
   return p && e.disposeIntermediateTensorInfo(d), x;
 }
-const he = {
-  kernelName: j,
+const me = {
+  kernelName: B,
   backendName: "webgl",
   kernelFunc: te
 };
-/**
- * @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.
- * =============================================================================
- */
 const ae = `
 if (a == b) {
   return 1.0;
@@ -524,16 +350,16 @@ return a / b;`, se = `
   }
 
   return result;
-`, ne = L({ opSnippet: ae, packedOpSnippet: se, checkOutOfBounds: !0 }), fe = {
-  kernelName: K,
+`, ne = L({ opSnippet: ae, packedOpSnippet: se, checkOutOfBounds: !0 }), xe = {
+  kernelName: F,
   backendName: "webgl",
   kernelFunc: ne
 };
 export {
   P as a,
-  he as b,
-  fe as c,
-  pe as d,
+  me as b,
+  xe as c,
+  fe as d,
   te as m,
   ne as r,
   Z as s,

package/dist/Reshape-CF6odzV4.js

@@ -0,0 +1,16 @@
+import "./index-ZyQhjEPo.js";
+import { s as p, n as m, a as d } from "./tensor-DdQUJZlz.js";
+import { b as c } from "./tensor_util-DV-FP5Q3.js";
+function i(t) {
+  const { inputs: h, attrs: o } = t, { x: e } = h, { shape: r } = o, a = p(e.shape), s = m(r, a), n = p(s);
+  return d(a === n, () => `The new shape (${s}) has ${n} elements and the old shape (${e.shape}) has ${a} elements. The new shape and old shape must have the same number of elements.`), t.backend.incRef(e.dataId), { dataId: e.dataId, shape: s, dtype: e.dtype };
+}
+const $ = {
+  kernelName: c,
+  backendName: "webgpu",
+  kernelFunc: i
+};
+export {
+  $ as a,
+  i as r
+};

package/dist/Reshape-_kILl6tK.js

@@ -0,0 +1,81 @@
+import "./index-ZyQhjEPo.js";
+import { u as C, g as f, a as R, b as g, c as I, d as c, e as u, i as m } from "./gpgpu_math-DDVJCn6-.js";
+import { b as x } from "./tensor_util-DV-FP5Q3.js";
+import { s as l, n as F, a as $ } from "./tensor-DdQUJZlz.js";
+class S {
+  constructor(t, i) {
+    this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0, this.customUniforms = [{ name: "inputShape", type: "ivec3" }], this.outputShape = t, this.enableShapeUniforms = C(this.outputShape.length);
+    let a = "";
+    for (let e = 0; e < 4; e++) {
+      let o = "thisRC = rc;";
+      e % 2 === 1 && (o += "thisRC.z += 1;"), e > 1 && (o += "thisRC.y += 1;"), a += `
+        ${o}
+        ${e > 0 ? "if(thisRC.y < rows && thisRC.z < cols){" : ""}
+          int flatIndex = getFlatIndex(thisRC);
+
+          ivec3 inputRC = inputCoordsFromReshapedOutCoords(flatIndex);
+          vec2 inputRCInnerDims = vec2(float(inputRC.y),float(inputRC.z));
+
+          result[${e}] =
+            getChannel(getA(inputRC.x, inputRC.y, inputRC.z), inputRCInnerDims);
+        ${e > 0 ? "}" : ""}
+      `;
+    }
+    this.userCode = `
+      ${b(i, this.enableShapeUniforms)}
+      ${this.enableShapeUniforms ? f() : R(t)}
+
+      void main() {
+        ivec3 rc = getOutputCoords();
+
+        vec4 result = vec4(0.);
+
+        ivec3 thisRC;
+        int rows = ${this.enableShapeUniforms ? "outShape[1]" : t[1]};
+        int cols = ${this.enableShapeUniforms ? "outShape[2]" : t[2]};
+
+        ${a}
+
+        setOutput(result);
+      }
+    `;
+  }
+}
+function b(s, t) {
+  return `
+    ivec3 inputCoordsFromReshapedOutCoords(int index) {
+      ${t ? g(["r", "c", "d"], "inputShape") : I(["r", "c", "d"], s)}
+      return ivec3(r, c, d);
+    }
+  `;
+}
+function v(s, t, i) {
+  const a = [
+    c(s.shape),
+    ...u(s.shape)
+  ], e = {
+    dtype: s.dtype,
+    shape: a,
+    dataId: s.dataId
+  }, o = [
+    c(t),
+    ...u(t)
+  ], r = new S(o, a), p = !0, n = [a], h = i.runWebGLProgram(r, [e], s.dtype, n, p);
+  return { dataId: h.dataId, shape: t, dtype: h.dtype };
+}
+function y(s) {
+  const { inputs: t, backend: i, attrs: a } = s, { x: e } = t, { shape: o } = a, r = i, p = l(e.shape), n = F(o, p), h = l(n);
+  $(p === h, () => `The new shape (${n}) has ${h} elements and the old shape (${e.shape}) has ${p} elements. The new shape and old shape must have the same number of elements.`);
+  const d = r.texData.get(e.dataId);
+  return d.isPacked && !m(e.shape, n) && !(d.texture !== null && m(d.shape, n)) ? v(e, n, r) : (r.incRef(e.dataId), { dataId: e.dataId, shape: n, dtype: e.dtype });
+}
+const O = {
+  kernelName: x,
+  backendName: "webgl",
+  kernelFunc: y
+};
+export {
+  S as R,
+  O as a,
+  y as r
+};

package/dist/TeachableLLM.js

@@ -2,45 +2,51 @@ import { defaultConfig as d } from "./models/config.js";
 import { saveModel as l } from "./loader/save.js";
 import { loadModel as _ } from "./loader/load.js";
 import u from "./Generator.js";
-import f from "./Trainer.js";
-import { E as p } from "./index-Dwqa6Zy2.js";
+import p from "./Trainer.js";
+import { E as f } from "./index-DvYrXKkX.js";
 import { dummyPassTrainAsync as m } from "./utilities/dummy.js";
-import "./index-BzFyqcy-.js";
+import "./utilities/packed.js";
+import "./index-ZyQhjEPo.js";
 import "./ops/cpu/attentionMask.js";
 import "./ops/webgl/attentionMask.js";
 import "./ops/grads/attentionMask.js";
-import "./ops/cpu/qkv.js";
-import "./ops/webgl/qkv.js";
-import "./ops/grads/qkv.js";
-import "./random_width-CXVRloNK.js";
-import "./register_all_kernels-DIGpEwcf.js";
-import "./index-Tf7vU29b.js";
-import "./dataset-DlZtKmBq.js";
+import "./random_width-DY6Kk2Dl.js";
+import "./register_all_kernels-Bwu1PTuU.js";
+import "./index-Cp39cXWe.js";
+import "./dataset-0xP8GjwI.js";
 import "./ops/cpu/rope.js";
 import "./ops/webgl/rope.js";
-import "./ops/grads/rope.js";
+import "./rope-B5UUMsPi.js";
 import "./ops/cpu/appendCache.js";
 import "./ops/webgl/appendCache.js";
-import "./ops/cpu/fusedSoftmax.js";
-import "./ops/webgl/fusedSoftmax.js";
-import "./ops/grads/fusedSoftmax.js";
-import "./ops/cpu/matMulGelu.js";
-import "./ops/webgl/matMulGelu.js";
-import "./ops/grads/matMulGelu.js";
+import "./ops/grads/softmax16.js";
+import "./matMul16--R5hOwDG.js";
+import "./ops/webgl/matMul16.js";
+import "./ops/cpu/matMul16.js";
+import "./pack16-CFUqumar.js";
+import "./ops/transpose16.js";
+import "./ops/reshape16.js";
+import "./ops/cpu/qkv.js";
+import "./ops/webgl/qkv.js";
+import "./ops/grads/qkv.js";
 import "./ops/cpu/normRMS.js";
 import "./ops/webgl/normRMS.js";
 import "./ops/grads/normRMS.js";
+import "./ops/grads/add16.js";
 import "./ops/cpu/gatherSub.js";
 import "./ops/webgl/gatherSub.js";
 import "./ops/cpu/scatterSub.js";
 import "./ops/webgl/scatterSub.js";
 import c from "./tokeniser/CharTokeniser.js";
 import g from "./tokeniser/bpe.js";
-import "./papaparse.min-C8l2Kvo1.js";
-import "./jszip.min-CjP2V1VV.js";
+import "./papaparse.min-C0cScC2i.js";
+import "./jszip.min-Bz5-11Bk.js";
+import "./ops/cpu/matMulGelu.js";
+import "./ops/webgl/matMulGelu.js";
+import "./ops/grads/matMulGelu.js";
 import "./ops/cpu/gelu.js";
 import "./ops/webgl/gelu.js";
-import "./gelu-Bp_-935b.js";
+import "./gelu-CNLFZWea.js";
 import "./ops/webgl/log.js";
 import "./ops/cpu/adamMoments.js";
 import "./ops/webgl/adamMoments.js";
@@ -51,7 +57,7 @@ import "./checks/normRMSGrad.js";
 import k from "./utilities/profile.js";
 import w from "./models/factory.js";
 class a {
-  ee = new p();
+  ee = new f();
   _config;
   _model;
   _tokeniser;
@@ -150,7 +156,7 @@ class a {
   trainer() {
     if (!this._model || !this._tokeniser)
       throw new Error("model_or_tokeniser_not_initialized.");
-    const t = new f(this._model, this._tokeniser);
+    const t = new p(this._model, this._tokeniser);
     return t.on("start", () => this.setStatus("training")), t.on("stop", () => this.setStatus("ready")), t.on("log", async (e, r) => {
       const o = this.ee.listeners("trainStep");
       for (const s of o)

package/dist/Trainer.d.ts

@@ -12,6 +12,8 @@ export interface ITrainerOptions {
     validationSplit?: number;
     advancedMetrics?: boolean;
     gradientCheckpointing?: boolean;
+    gradientMetrics?: boolean;
+    mixedPrecision?: boolean;
 }
 interface ExtendedTrainingProgress extends TrainingProgress {
     progress: number;

package/dist/Trainer.js

@@ -1,4 +1,4 @@
-import { E as l } from "./index-Dwqa6Zy2.js";
+import { E as l } from "./index-DvYrXKkX.js";
 import h from "./training/FullTrainer.js";
 class m extends l {
   trainer;
@@ -28,7 +28,7 @@ class m extends l {
   async train(t) {
     if (!this.trainDataset || !this.validationDataset)
       throw new Error("Datasets not prepared");
-    this.hasTrained || this.trainer.setLearningRate(t?.learningRate || 1e-3), this.hasTrained = !0, this.emit("start"), this.trainer.setGradientCheckpointing(t?.gradientCheckpointing || !1), await this.trainer.trainOnDataset(
+    this.hasTrained || this.trainer.setLearningRate(t?.learningRate || 1e-3), this.hasTrained = !0, this.emit("start"), this.trainer.setGradientCheckpointing(t?.gradientCheckpointing || !1), this.trainer.setMixedPrecision(t?.mixedPrecision || !1), await this.trainer.trainOnDataset(
       this.trainDataset,
       {
         prompt: t?.prompt,
@@ -36,6 +36,7 @@ class m extends l {
         desiredLoss: t?.desiredLoss || 0.01,
         maxSteps: t?.maxSteps || 1e3,
         advancedMetrics: t?.advancedMetrics || !1,
+        gradientMetrics: t?.gradientMetrics || !1,
         onStep: async (e, a) => {
           this.log.push(e), this.progress = {
             ...a,

package/dist/{axis_util-TbGYJ208.js → axis_util-BvHEw88j.js}

@@ -1,20 +1,4 @@
-import { n as c } from "./index-BzFyqcy-.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.
- * =============================================================================
- */
+import { a as c } from "./tensor-DdQUJZlz.js";
 function i(e, n) {
   for (let t = 0; t < e.length; ++t)
     if (e[e.length - t - 1] !== n - 1 - t)
@@ -60,12 +44,12 @@ function x(e, n) {
   return t;
 }
 export {
-  x as a,
-  m as b,
+  d as a,
+  x as b,
   l as c,
-  i as d,
+  m as d,
   h as e,
-  a as f,
-  d as g,
-  g as h
+  i as f,
+  g,
+  a as h
 };

package/dist/backend.d.ts

@@ -1 +1,2 @@
-export declare function selectBackend(backendName: 'cpu' | 'webgl' | 'webgpu'): Promise<void>;
+import { GPUOptions } from './patches/webgpu_base';
+export declare function selectBackend(backendName: 'cpu' | 'webgl' | 'webgpu', options?: GPUOptions): Promise<void>;

package/dist/backend.js

@@ -1,7 +1,13 @@
-import { g as a, s as i, r as o } from "./index-BzFyqcy-.js";
-async function e(t) {
-  a() !== t && (t === "webgpu" && (await import("./index-C1rx_Ajs.js"), await import("./ops/webgpu/index.js")), await i(t), await o(), console.log(`Backend set to ${t}`));
+import { g as o, s as e, r as s } from "./index-ZyQhjEPo.js";
+async function c(t, a) {
+  if (o() !== t) {
+    if (t === "webgpu") {
+      const { registerWebGPUBackend: i } = await import("./patches/webgpu_base.js");
+      i(a), await import("./index-CjOj7j-u.js"), await import("./ops/webgpu/index.js");
+    }
+    await e(t), await s(), console.log(`Backend set to ${t}`);
+  }
 }
 export {
-  e as selectBackend
+  c as selectBackend
 };