@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/inference/types.d.ts

@@ -0,0 +1,16 @@
+import { Conversation } from '../tokeniser/type';
+export interface GeneratorConversation extends Conversation {
+    _completed?: boolean;
+    _timestamp?: number;
+}
+export interface GenerateOptions {
+    temperature?: number;
+    topK?: number;
+    topP?: number;
+    usePadding?: boolean;
+    attentionScores?: boolean;
+    includeProbabilities?: boolean;
+    embeddings?: 'embedding' | 'logits' | 'softmax' | 'all';
+    targets?: number[];
+    loraName?: string;
+}

package/dist/kernel_funcs_utils-HiXOOx3f.js

@@ -0,0 +1,229 @@
+import { Ci as e, Di as t, In as n, Ms as r, Ya as i, ca as a, ko as o, oc as s, to as c } from "./dist-BewPQWjc.js";
+import { r as l } from "./backend_util-Cg-roD1p.js";
+import { a as u, o as d } from "./gpgpu_math-DvLcCH6u.js";
+//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/packing_util.js
+function f(e, t) {
+	return [
+		"x",
+		"y",
+		"z",
+		"w",
+		"u",
+		"v"
+	].slice(0, t).map((t) => `${e}.${t}`);
+}
+function p(e, t) {
+	return t === 1 ? [e] : f(e, t);
+}
+function m(e, t) {
+	if (e === 1) return "rc";
+	let n = "";
+	for (let r = 0; r < e; r++) n += t[r], r < e - 1 && (n += ",");
+	return n;
+}
+//#endregion
+//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/unaryop_gpu.js
+var h = class {
+	constructor(e, t) {
+		this.variableNames = ["A"], this.outputShape = e, this.enableShapeUniforms = u(this.outputShape.length), this.userCode = `
+      float unaryOperation(float x) {
+        ${t}
+      }
+
+      void main() {
+        float x = getAAtOutCoords();
+        float y = unaryOperation(x);
+
+        setOutput(y);
+      }
+    `;
+	}
+}, g = "if (isnan(x)) return x;", _ = "return x;", v = "return abs(x);", y = "return (x >= 0.0) ? x : (exp(x) - 1.0);", b = g + "\n  return (x < 0.0) ? 0.0 : x;\n", x = g + "\n  return (x < 0.0) ? 0.0 : min(6.0, x);\n", S = "return x;", C = "return 1.0 / (1.0 + exp(-1.0 * x));", w = "return x;", T = "\n  vec4 result;\n\n  result.r = (x.r >= 0.0) ? x.r : (exp(x.r) - 1.0);\n  result.g = (x.g >= 0.0) ? x.g : (exp(x.g) - 1.0);\n  result.b = (x.b >= 0.0) ? x.b : (exp(x.b) - 1.0);\n  result.a = (x.a >= 0.0) ? x.a : (exp(x.a) - 1.0);\n\n  return result;\n", E = "\n  vec4 result = x * vec4(greaterThanEqual(x, vec4(0.0)));\n  bvec4 isNaN = isnan(x);\n\n  result.r = isNaN.r ? x.r : result.r;\n  result.g = isNaN.g ? x.g : result.g;\n  result.b = isNaN.b ? x.b : result.b;\n  result.a = isNaN.a ? x.a : result.a;\n\n  return result;\n", D = "\n  vec4 result = min(x, vec4(6.)) * vec4(greaterThanEqual(x, vec4(0.0)));\n  bvec4 isNaN = isnan(x);\n\n  result.r = isNaN.r ? x.r : result.r;\n  result.g = isNaN.g ? x.g : result.g;\n  result.b = isNaN.b ? x.b : result.b;\n  result.a = isNaN.a ? x.a : result.a;\n\n  return result;\n", O = "return 1.0 / (1.0 + exp(-1.0 * x));", k = class {
+	constructor(e, t) {
+		this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0, this.outputShape = e, this.enableShapeUniforms = u(this.outputShape.length), this.userCode = `
+      vec4 unaryOperation(vec4 x) {
+        ${t}
+      }
+
+      void main() {
+        vec4 x = getAAtOutCoords();
+        vec4 y = unaryOperation(x);
+
+        setOutput(y);
+      }
+    `;
+	}
+}, A = "\n  if (isnan(a)) return a;\n  if (isnan(b)) return b;\n", j = class {
+	constructor(e, t, r) {
+		this.variableNames = ["A", "B"], this.outputShape = n(t, r), this.enableShapeUniforms = u(this.outputShape.length), this.userCode = `
+      float binaryOperation(float a, float b) {
+        ${e}
+      }
+
+      void main() {
+        float a = getAAtOutCoords();
+        float b = getBAtOutCoords();
+        setOutput(binaryOperation(a, b));
+      }
+    `;
+	}
+}, M = "\n  result.r = isNaN.r ? NAN : result.r;\n  result.g = isNaN.g ? NAN : result.g;\n  result.b = isNaN.b ? NAN : result.b;\n  result.a = isNaN.a ? NAN : result.a;\n", N = class {
+	constructor(e, t, r, i = !1) {
+		this.variableNames = ["A", "B"], this.supportsBroadcasting = !0, this.packedInputs = !0, this.packedOutput = !0, this.outputShape = n(t, r);
+		let a = this.outputShape.length;
+		this.enableShapeUniforms = u(a);
+		let o = "";
+		if (i) if (a === 0 || s(this.outputShape) === 1) o = "\n          result.y = 0.;\n          result.z = 0.;\n          result.w = 0.;\n        ";
+		else if (o = `
+          ${d(a)} coords = getOutputCoords();
+        `, a === 1) this.enableShapeUniforms ? o += "\n            result.y = (coords + 1) >= outShape ? 0. : result.y;\n            result.z = 0.;\n            result.w = 0.;\n          " : o += `
+            result.y = (coords + 1) >= ${this.outputShape[0]} ? 0. : result.y;
+            result.z = 0.;
+            result.w = 0.;
+          `;
+		else {
+			let e = p("coords", a);
+			this.enableShapeUniforms ? o += `
+            bool nextRowOutOfBounds =
+              (${e[a - 2]} + 1) >= outShape[${a} - 2];
+            bool nextColOutOfBounds =
+              (${e[a - 1]} + 1) >= outShape[${a} - 1];
+            result.y = nextColOutOfBounds ? 0. : result.y;
+            result.z = nextRowOutOfBounds ? 0. : result.z;
+            result.w = nextColOutOfBounds || nextRowOutOfBounds ? 0. : result.w;
+          ` : o += `
+            bool nextRowOutOfBounds =
+              (${e[a - 2]} + 1) >= ${this.outputShape[a - 2]};
+            bool nextColOutOfBounds =
+              (${e[a - 1]} + 1) >= ${this.outputShape[a - 1]};
+            result.y = nextColOutOfBounds ? 0. : result.y;
+            result.z = nextRowOutOfBounds ? 0. : result.z;
+            result.w = nextColOutOfBounds || nextRowOutOfBounds ? 0. : result.w;
+          `;
+		}
+		this.userCode = `
+      vec4 binaryOperation(vec4 a, vec4 b) {
+        ${e}
+      }
+
+      void main() {
+        vec4 a = getAAtOutCoords();
+        vec4 b = getBAtOutCoords();
+
+        vec4 result = binaryOperation(a, b);
+        ${o}
+
+        setOutput(result);
+      }
+    `;
+	}
+};
+//#endregion
+//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/kernels/Identity.js
+function P(e) {
+	let { inputs: t, backend: n } = e, { x: r } = t;
+	return n.incRef(r.dataId), {
+		dataId: r.dataId,
+		shape: r.shape,
+		dtype: r.dtype
+	};
+}
+var F = {
+	kernelName: i,
+	backendName: "webgl",
+	kernelFunc: P
+};
+//#endregion
+//#region node_modules/@tensorflow/tfjs-backend-webgl/dist/kernels/Complex.js
+function I(e) {
+	let { inputs: t, backend: n } = e, { real: r, imag: i } = t, a = n.makeTensorInfo(r.shape, "complex64"), o = n.texData.get(a.dataId);
+	return o.complexTensorInfos = {
+		real: P({
+			inputs: { x: r },
+			backend: n
+		}),
+		imag: P({
+			inputs: { x: i },
+			backend: n
+		})
+	}, a;
+}
+var L = {
+	kernelName: a,
+	backendName: "webgl",
+	kernelFunc: I
+}, R = "return (a < 0.) ? b * a : a;", z = "\n  vec4 aLessThanZero = vec4(lessThan(a, vec4(0.)));\n  return (aLessThanZero * (b * a)) + ((vec4(1.0) - aLessThanZero) * a);\n";
+function B(e) {
+	let { inputs: n, backend: i, attrs: a } = e, { x: o } = n, { alpha: s } = a, c = i.makeTensorInfo([], "float32", t(s, "float32")), l = r().getBool("WEBGL_PACK_BINARY_OPERATIONS") ? new N(z, o.shape, c.shape) : new j(R, o.shape, c.shape), u = i.runWebGLProgram(l, [o, c], "float32");
+	return i.disposeIntermediateTensorInfo(c), u;
+}
+var V = {
+	kernelName: c,
+	backendName: "webgl",
+	kernelFunc: B
+}, H = "return (a < 0.) ? b * a : a;", U = "\n  vec4 aLessThanZero = vec4(lessThan(a, vec4(0.)));\n  return (aLessThanZero * (b * a)) + ((vec4(1.0) - aLessThanZero) * a);\n";
+function W(e) {
+	let { inputs: t, backend: n } = e, { x: i, alpha: a } = t, o = r().getBool("WEBGL_PACK_BINARY_OPERATIONS") ? new N(U, i.shape, a.shape) : new j(H, i.shape, a.shape);
+	return n.runWebGLProgram(o, [i, a], "float32");
+}
+var G = {
+	kernelName: o,
+	backendName: "webgl",
+	kernelFunc: W
+}, K = "if (isnan(x)) return x;";
+function q({ opSnippet: e, packedOpSnippet: t, cpuKernelImpl: n, dtype: i }) {
+	return ({ inputs: a, backend: o }) => {
+		let { x: s } = a, c = o, l = i || s.dtype;
+		if (c.shouldExecuteOnCPU([s]) && n != null) {
+			let e = n(c.texData.get(s.dataId).values, l);
+			return c.makeTensorInfo(s.shape, l, e);
+		}
+		let u = r().getBool("WEBGL_PACK_UNARY_OPERATIONS") && t != null, d;
+		return d = u ? new k(s.shape, t) : new h(s.shape, e), c.runWebGLProgram(d, [s], l);
+	};
+}
+function J({ opSnippet: t, packedOpSnippet: n, checkOutOfBounds: i = !1, supportsComplex: a = !1, cpuKernelImpl: o, dtype: s }) {
+	return ({ inputs: c, backend: u }) => {
+		let { a: d, b: f } = c, p = u;
+		if (a && d.dtype === "complex64") {
+			let n = p.texData.get(d.dataId), r = p.texData.get(f.dataId), [i, a] = [[n.complexTensorInfos.real, r.complexTensorInfos.real], [n.complexTensorInfos.imag, r.complexTensorInfos.imag]].map((n) => {
+				let [r, i] = n, a = {
+					dataId: r.dataId,
+					dtype: r.dtype,
+					shape: d.shape
+				}, o = {
+					dataId: i.dataId,
+					dtype: i.dtype,
+					shape: f.shape
+				}, s = new j(t, d.shape, f.shape);
+				return p.runWebGLProgram(s, [a, o], e(r.dtype, i.dtype));
+			}), o = I({
+				inputs: {
+					real: i,
+					imag: a
+				},
+				backend: p
+			});
+			return p.disposeIntermediateTensorInfo(i), p.disposeIntermediateTensorInfo(a), o;
+		}
+		let m = s || e(d.dtype, f.dtype);
+		if ((d.dtype === "string" || f.dtype === "string" || p.shouldExecuteOnCPU([d, f])) && o != null) {
+			let e = p.texData.get(d.dataId).values, t = p.texData.get(f.dataId).values, n = d.dtype === "string" ? l(e) : e, r = d.dtype === "string" ? l(t) : t, [i, a] = o(d.shape, f.shape, n, r, m), s = p.makeTensorInfo(a, m), c = p.texData.get(s.dataId);
+			return c.values = i, s;
+		}
+		let h = r().getBool("WEBGL_PACK_BINARY_OPERATIONS") && n != null, g;
+		return g = h ? new N(n, d.shape, f.shape, i) : new j(t, d.shape, f.shape), p.runWebGLProgram(g, [d, f], m);
+	};
+}
+function Y(e, t = !1) {
+	if (e === "linear") return t ? w : _;
+	if (e === "relu") return t ? E : b;
+	if (e === "elu") return t ? T : y;
+	if (e === "relu6") return t ? D : x;
+	if (e === "prelu") return t ? U : H;
+	if (e === "leakyrelu") return t ? z : R;
+	if (e === "sigmoid") return t ? O : C;
+	throw Error(`Activation ${e} has not been implemented for the WebGL backend.`);
+}
+//#endregion
+export { f as S, g as _, G as a, p as b, L as c, N as d, M as f, v as g, k as h, q as i, P as l, A as m, J as n, V as o, j as p, Y as r, I as s, K as t, F as u, S as v, m as x, h as y };

package/dist/layers/BaseLayer.d.ts

@@ -0,0 +1,44 @@
+import { GPTConfig } from '../../models/config';
+import { default as MemoryProfiler } from '../../utilities/profile';
+import { default as RoPECache } from './RoPECache';
+import { Tensor, Variable } from '@tensorflow/tfjs-core';
+import { default as WeightStore } from './WeightStore';
+export interface ForwardAttributes {
+    training: boolean;
+    checkpointing?: boolean;
+    mixedPrecision?: boolean;
+    ropeCache?: RoPECache;
+    outputEmbeddings?: boolean;
+    embeddings?: {
+        name: string;
+        tensor: Tensor;
+    }[];
+    dropout?: number;
+    layerDrop?: number;
+}
+export default abstract class BaseLayer<ATTR extends ForwardAttributes = ForwardAttributes, CONFIG extends GPTConfig = GPTConfig> {
+    readonly parent?: BaseLayer;
+    readonly config: CONFIG;
+    weightStore: WeightStore;
+    readonly children: BaseLayer[];
+    private profiler?;
+    private ownVariables;
+    constructor(config: CONFIG, parent?: BaseLayer);
+    getProfiler(): MemoryProfiler | undefined;
+    setProfiler(profiler: MemoryProfiler | null): void;
+    startMemory(): void;
+    endMemory(label: string): void;
+    addVariable(name: string, variable?: Variable): void;
+    addChildVariable(name: string): void;
+    get variables(): Variable[];
+    get trainableVariables(): Variable[];
+    getVariable(name: string): Tensor;
+    hasVariable(name: string): boolean;
+    setVariable(name: string, variable: Variable): void;
+    dispose(): void;
+    protected build(): void;
+    abstract forward(attrs: ATTR, ...x: Tensor[]): Tensor | Tensor[];
+    call(attrs: ATTR, ...x: Tensor[]): Tensor | Tensor[];
+    callCheckpoint(attrs: ATTR, ...x: Tensor[]): Tensor;
+    private checkpointingFn;
+}

package/dist/layers/BaseLayer.js

@@ -0,0 +1,76 @@
+import { Gt as e, Wt as t, ii as n } from "../dist-BewPQWjc.js";
+import r from "./WeightStore.js";
+//#region lib/layers/BaseLayer.ts
+var i = class {
+	parent;
+	config;
+	weightStore;
+	children = [];
+	profiler;
+	ownVariables = /* @__PURE__ */ new Set();
+	constructor(e, t) {
+		this.config = e, this.parent = t, this.parent ? (this.parent.children.push(this), this.weightStore = this.parent.weightStore) : this.weightStore = new r();
+	}
+	getProfiler() {
+		return this.profiler;
+	}
+	setProfiler(e) {
+		this.profiler = e || void 0, this.children.forEach((t) => {
+			t.setProfiler(e);
+		});
+	}
+	startMemory() {
+		this.profiler?.startMemory();
+	}
+	endMemory(e) {
+		this.profiler?.endMemory(e);
+	}
+	addVariable(e, t) {
+		this.weightStore.addVariable(e, t), this.ownVariables.add(e), this.parent && this.parent.addChildVariable(e);
+	}
+	addChildVariable(e) {
+		this.ownVariables.add(e);
+	}
+	get variables() {
+		return this.weightStore.variables;
+	}
+	get trainableVariables() {
+		return this.weightStore.trainableVariables.filter((e) => this.ownVariables.has(e.name));
+	}
+	getVariable(e) {
+		return this.weightStore.getVariable(e);
+	}
+	hasVariable(e) {
+		return this.weightStore.hasVariable(e);
+	}
+	setVariable(e, t) {
+		this.weightStore.setVariable(e, t);
+	}
+	dispose() {
+		this.weightStore.dispose();
+	}
+	build() {}
+	call(e, ...t) {
+		return this.build(), this.forward(e, ...t);
+	}
+	callCheckpoint(e, ...t) {
+		return this.build(), this.checkpointingFn(e, ...t);
+	}
+	checkpointingFn(r, ...i) {
+		let a = this.trainableVariables;
+		return t((...t) => {
+			let o = t[t.length - 1], s = t.slice(0, i.length), c = this.forward(r, ...s);
+			return o(s), {
+				value: c,
+				gradFunc: (t, i) => {
+					let o = n().state.activeTape;
+					n().state.activeTape = [];
+					let c = e((...e) => this.forward(r, ...e.slice(0, s.length)))([...i, ...a], t);
+					return n().state.activeTape = o, c;
+				}
+			};
+		})(...i, ...a);
+	}
+};
+//#endregion
+export { i as default };

package/dist/layers/CausalSelfAttention.d.ts

@@ -0,0 +1,39 @@
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { Tensor } from '@tensorflow/tfjs-core';
+import { GPTConfig } from '../../models/config';
+export interface KVCache {
+    k?: Tensor;
+    v?: Tensor;
+    length: number;
+    cumulativeLength: number;
+}
+export interface AttentionScores {
+    meanOfHeads?: boolean;
+    attentionOut?: Tensor[];
+}
+interface AttentionForwardAttributes extends ForwardAttributes {
+    attentionScores?: AttentionScores;
+    pastKV?: KVCache;
+    seed?: number;
+    ropePositionOffset?: number;
+}
+export interface CausalSelfAttentionConfig {
+    useQKNorm?: boolean;
+}
+export default class CausalSelfAttention extends BaseLayer<AttentionForwardAttributes> {
+    private readonly attentionConfig;
+    private divisor;
+    private index;
+    private units;
+    private projUnits;
+    private ATTN;
+    private PROJ;
+    constructor(index: number, config: GPTConfig, attentionConfig: CausalSelfAttentionConfig, parent?: BaseLayer);
+    protected build(): void;
+    private getAttentionScores;
+    private getQKV;
+    private getOutputProjection;
+    private updateCache;
+    forward(attr: AttentionForwardAttributes, x: Tensor): Tensor;
+}
+export {};

package/dist/layers/CausalSelfAttention.js

@@ -0,0 +1,99 @@
+import { T as e, di as t, oi as n, pt as r } from "../dist-BewPQWjc.js";
+import { isPackedTensor as i } from "../utilities/packed.js";
+import { transpose16 as a } from "../ops/transpose16.js";
+import { reshape16 as o } from "../ops/reshape16.js";
+import { t as s } from "../matMul16-BNfZSnNM.js";
+import { t as c } from "../pack16-Ck-spx_F.js";
+import { attentionMask as l } from "../ops/attentionMask.js";
+import u from "./BaseLayer.js";
+import { t as d } from "../rope-CbeGlsV8.js";
+import { appendCache as f } from "../ops/appendCache.js";
+import { softmax16 as p } from "../ops/softmax16.js";
+import { dot16 as m } from "../ops/dot16.js";
+import { qkv as h } from "../ops/qkv.js";
+import { normRMS as g } from "../ops/normRMS.js";
+import { dropout16 as _ } from "../ops/dropout16.js";
+//#region lib/layers/CausalSelfAttention.ts
+var v = class extends u {
+	attentionConfig;
+	divisor;
+	index;
+	units;
+	projUnits;
+	ATTN;
+	PROJ;
+	constructor(e, t, n, r) {
+		super(t, r), this.attentionConfig = n, this.index = e, this.units = t.nEmbed * 3, this.projUnits = t.nEmbed, this.ATTN = `block_${this.index}_cAttn`, this.PROJ = `block_${this.index}_cProj`, this.addVariable(this.ATTN), this.addVariable(this.PROJ), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead);
+	}
+	build() {
+		this.hasVariable(this.ATTN) === !1 && this.setVariable(this.ATTN, e(r([this.config.nEmbed, this.units], 0, .02), !0, this.ATTN)), this.hasVariable(this.PROJ) === !1 && this.setVariable(this.PROJ, e(r([this.projUnits, this.config.nEmbed], 0, .02), !0, this.PROJ));
+	}
+	getAttentionScores(e, t, n) {
+		let r = l(e, t, this.divisor, n), i = p(r);
+		return r.dispose(), i;
+	}
+	getQKV(e) {
+		let t = i(e) ? c(this.getVariable(this.ATTN)) : this.getVariable(this.ATTN), n = h(e, t, this.config.nHead);
+		return i(e) && t.dispose(), n;
+	}
+	getOutputProjection(e) {
+		let t = e.shape[0], n = e.shape[2], r = this.config.nEmbed, s = i(e), l = a(e, [
+			0,
+			2,
+			1,
+			3
+		]), u = o(l, [
+			t,
+			n,
+			s ? r / 2 : r
+		]);
+		l.dispose();
+		let d = s ? c(this.getVariable(this.PROJ)) : this.getVariable(this.PROJ), f = m(u, d);
+		return s && d.dispose(), u.dispose(), f;
+	}
+	updateCache(e, t, r) {
+		let i = this.config.blockSize, a = e.shape[2], o = r.length || 0, s = f(e, i, o, r.k);
+		e.dispose(), r.k && r.k.dispose();
+		let c = f(t, i, o, r.v);
+		t.dispose(), r.v && r.v.dispose();
+		let l = Math.min(o + a, i), u = r.cumulativeLength + a;
+		r.length = l, r.cumulativeLength = u, r.k = n(s), r.v = n(c);
+	}
+	forward(e, r) {
+		return t(() => {
+			this.startMemory();
+			let [t, i, a] = this.getQKV(r), o = e.pastKV ? e.pastKV.cumulativeLength : e.ropePositionOffset || 0, c = e.ropeCache, l = c ? d(t, c, o) : t, u = c ? d(i, c, o) : i, f = this.attentionConfig.useQKNorm ?? !1, p = f ? g(l) : l;
+			f && l.dispose();
+			let m = f ? g(u) : u;
+			f && u.dispose(), c && (t.dispose(), i.dispose());
+			let h = e.pastKV ? e.pastKV.length : 0;
+			e.pastKV && !e.training && this.updateCache(m, a, e.pastKV);
+			let v = e.pastKV?.k ? e.pastKV.k : m, y = e.pastKV?.v ? e.pastKV.v : a, b;
+			b = h > 0 ? this.getAttentionScores(p, v, h) : this.getAttentionScores(p, v), p.dispose(), e.pastKV || v.dispose();
+			let x = s(b, y), S = e.attentionScores !== void 0 && e.attentionScores.attentionOut !== void 0;
+			S || b.dispose(), e.pastKV || y.dispose();
+			let C = this.getOutputProjection(x);
+			if (x.dispose(), S && e.attentionScores && e.attentionScores.attentionOut !== void 0) {
+				let t = b.shape[1], r = b.shape[2];
+				e.attentionScores.attentionOut?.push(n(b.slice([
+					0,
+					0,
+					0,
+					0
+				], [
+					1,
+					-1,
+					-1,
+					-1
+				]).reshape([
+					t,
+					r,
+					-1
+				])));
+			}
+			return this.endMemory("CausalSelfAttention"), e.dropout && e.dropout > 0 ? _(C, e.dropout) : C;
+		});
+	}
+};
+//#endregion
+export { v as default };

package/dist/layers/LoRA.d.ts

@@ -0,0 +1,14 @@
+import { default as WeightStore } from './WeightStore';
+export default class LoRA {
+    private weightStore;
+    readonly alpha: number;
+    readonly rank: number;
+    readonly variables: Set<string>;
+    private scale;
+    readonly name: string;
+    constructor(name: string, weightStore: WeightStore, alpha: number, rank: number, variables: string[]);
+    attach(): void;
+    merge(): void;
+    detach(): void;
+    dispose(): void;
+}

package/dist/layers/LoRA.js

@@ -0,0 +1,48 @@
+import { a as e } from "../chunk-BPntVaq0.js";
+import { T as t, _n as n, di as r, kt as i, pt as a } from "../dist-BewPQWjc.js";
+import { t as o } from "../picomatch-3tUnMMbd.js";
+//#region lib/layers/LoRA.ts
+var s = /* @__PURE__ */ e(o(), 1), c = class {
+	weightStore;
+	alpha;
+	rank;
+	variables;
+	scale;
+	name;
+	constructor(e, r, o, c, l) {
+		this.name = e, this.weightStore = r, this.alpha = o, this.rank = c;
+		let u = (0, s.default)(l), d = r.variableNames.filter((e) => u(e) && !e.endsWith("_loraA") && !e.endsWith("_loraB"));
+		this.variables = new Set(d), this.scale = n(o / c), this.variables.forEach((e) => {
+			let n = this.weightStore.getRawVariable(e), [r, o] = n.shape, s = `${e}_${this.name}_loraA`, c = `${e}_${this.name}_loraB`;
+			if (n.shape.length !== 2) {
+				console.warn(`LoRA currently only supports 2D weight matrices. Variable ${e} has shape ${n.shape}`), this.variables.delete(e);
+				return;
+			}
+			this.weightStore.hasVariable(s) || this.weightStore.hasVariable(c) || (this.weightStore.addVariable(s, t(a([r, this.rank], 0, .02), !0, s)), this.weightStore.addVariable(c, t(i([this.rank, o]), !0, c)));
+		});
+	}
+	attach() {
+		if (this.weightStore.onWeightRead) throw Error("LoRA cannot be applied to a WeightStore that already has a onWeightRead hook.");
+		this.weightStore.onWeightRead = (e, t) => this.variables.has(e) ? r(() => {
+			let n = this.weightStore.getRawVariable(`${e}_${this.name}_loraA`), r = this.weightStore.getRawVariable(`${e}_${this.name}_loraB`);
+			return t.add(n.matMul(r).mul(this.scale));
+		}) : t, this.weightStore.setTrainable([`*_${this.name}_loraA`, `*_${this.name}_loraB`]);
+	}
+	merge() {
+		this.variables.forEach((e) => {
+			let t = this.weightStore.getRawVariable(e), n = this.weightStore.getRawVariable(`${e}_${this.name}_loraA`), i = this.weightStore.getRawVariable(`${e}_${this.name}_loraB`), a = r(() => t.add(n.matMul(i).mul(this.scale)));
+			t.assign(a), a.dispose();
+		});
+	}
+	detach() {
+		this.weightStore.onWeightRead = void 0, this.weightStore.setTrainable(["*"]);
+	}
+	dispose() {
+		this.detach(), this.scale.dispose(), this.variables.forEach((e) => {
+			let t = `${e}_${this.name}_loraA`, n = `${e}_${this.name}_loraB`;
+			this.weightStore.getRawVariable(t).dispose(), this.weightStore.getRawVariable(n).dispose(), this.weightStore.deleteVariable(t), this.weightStore.deleteVariable(n);
+		}), this.variables.clear();
+	}
+};
+//#endregion
+export { c as default };

package/dist/layers/MLP.d.ts

@@ -0,0 +1,17 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../../main';
+export interface MLPConfig {
+    activation?: 'gelu' | 'relu2';
+    hiddenFactor?: number;
+}
+export default class MLP extends BaseLayer {
+    private index;
+    private hiddenUnits;
+    private MLPHIDDEN;
+    private MLPOUT;
+    private mlpConfig;
+    constructor(index: number, config: GPTConfig, mlpConfig: MLPConfig, parent?: BaseLayer);
+    protected build(): void;
+    forward(attr: ForwardAttributes, x: Tensor): Tensor;
+}

package/dist/layers/MLP.js

@@ -0,0 +1,34 @@
+import { T as e, di as t, pt as n } from "../dist-BewPQWjc.js";
+import { reshape16 as r } from "../ops/reshape16.js";
+import { t as i } from "../matMul16-BNfZSnNM.js";
+import a from "./BaseLayer.js";
+import { dropout16 as o } from "../ops/dropout16.js";
+//#region lib/layers/MLP.ts
+var s = class extends a {
+	index;
+	hiddenUnits;
+	MLPHIDDEN;
+	MLPOUT;
+	mlpConfig;
+	constructor(e, t, n, r) {
+		super(t, r), this.index = e, this.mlpConfig = n, this.hiddenUnits = (n.hiddenFactor ?? t.mlpFactor) * t.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
+	}
+	build() {
+		this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(this.MLPHIDDEN, e(n([this.config.nEmbed, this.hiddenUnits], 0, .02), !0, this.MLPHIDDEN)), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(this.MLPOUT, e(n([this.hiddenUnits, this.config.nEmbed], 0, .02 / Math.sqrt(2 * this.config.nLayer)), !0, this.MLPOUT));
+	}
+	forward(e, n) {
+		return t(() => {
+			this.startMemory();
+			let [t, a, s] = n.shape, c = i(r(n, [t * a, s]), this.getVariable(this.MLPHIDDEN), !1, !1, { activation: this.mlpConfig.activation ?? "gelu" }), l = i(c, this.getVariable(this.MLPOUT));
+			c.dispose();
+			let u = r(l, [
+				t,
+				a,
+				s
+			]);
+			return this.endMemory("MLP"), e.dropout && e.dropout > 0 ? o(u, e.dropout) : u;
+		});
+	}
+};
+//#endregion
+export { s as default };

package/dist/layers/PositionEmbedding.d.ts

@@ -0,0 +1,8 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer } from './BaseLayer';
+import { GPTConfig, ModelForwardAttributes } from '../../main';
+export default class PositionEmbedding extends BaseLayer {
+    private wpe?;
+    constructor(config: GPTConfig, name?: string, parent?: BaseLayer);
+    forward(attrs: ModelForwardAttributes, x: Tensor): Tensor;
+}

package/dist/layers/PositionEmbedding.js

@@ -0,0 +1,27 @@
+import { Tt as e, _n as t, di as n, dt as r, vi as i } from "../dist-BewPQWjc.js";
+import { n as a, t as o } from "../dist-VEU5mfO0.js";
+import s from "./BaseLayer.js";
+//#region lib/layers/PositionEmbedding.ts
+var c = class extends s {
+	wpe;
+	constructor(e, t = "", n) {
+		super(e, n), this.wpe = o({
+			inputDim: this.config.blockSize,
+			outputDim: this.config.nEmbed,
+			name: t,
+			embeddingsInitializer: a({
+				mean: 0,
+				stddev: .02
+			})
+		});
+	}
+	forward(a, o) {
+		let s = a.cache?.[0]?.length ?? 0;
+		return n(() => {
+			let [, n] = o.shape, a = this.config.blockSize, c = e(i(r(0, n, 1, "int32"), t(s, "int32")), t(a, "int32")), l = this.wpe.apply(c);
+			return o.add(l);
+		});
+	}
+};
+//#endregion
+export { c as default };

package/dist/layers/RMSNorm.d.ts

@@ -0,0 +1,12 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+import { default as BaseLayer, ForwardAttributes } from './BaseLayer';
+import { GPTConfig } from '../../main';
+export interface RMSNormConfig {
+    useGamma?: boolean;
+}
+export default class RMSNorm extends BaseLayer {
+    private GAMMA;
+    private rmsConfig;
+    constructor(config: GPTConfig, rmsConfig: RMSNormConfig, name?: string, parent?: BaseLayer);
+    forward(_: ForwardAttributes, x: Tensor): Tensor;
+}