@genai-fi/nanogpt 0.20.0 → 0.20.2

package/dist/ops/cpu/appendCache.js

@@ -0,0 +1,65 @@
+import { Ii as e, _r as t } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/appendCache.ts
+function n(e) {
+	let { cache: n, item: r } = e.inputs, { maxSize: i, pastLen: a } = e.attrs, o = n.shape[0], s = n.shape[1], c = n.shape[3], l = r.shape[2];
+	if (a + l <= i) {
+		let e = n.slice([
+			0,
+			0,
+			0,
+			0
+		], [
+			o,
+			s,
+			a,
+			c
+		]), u = n.slice([
+			0,
+			0,
+			a + l,
+			0
+		], [
+			o,
+			s,
+			i - a - l,
+			c
+		]), d = l < l ? r.slice([
+			0,
+			0,
+			0,
+			0
+		], [
+			o,
+			s,
+			l,
+			c
+		]) : r, f = t([
+			e,
+			d,
+			u
+		], 2);
+		return e.dispose(), u.dispose(), d !== r && d.dispose(), f;
+	}
+	let u = n.slice([
+		0,
+		0,
+		l,
+		0
+	], [
+		o,
+		s,
+		i - l,
+		c
+	]), d = t([u, r], 2);
+	return u.dispose(), d;
+}
+e({
+	kernelName: "AppendCache",
+	backendName: "cpu",
+	kernelFunc: n
+}), e({
+	kernelName: "AppendCache",
+	backendName: "tensorflow",
+	kernelFunc: n
+});
+//#endregion

package/dist/ops/cpu/attentionMask.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/attentionMask.js

@@ -0,0 +1,16 @@
+import { Ii as e, Ot as t, Pn as n, _n as r, ar as i, gr as a, kt as o, p as s } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/attentionMask.ts
+function c(e) {
+	let { q: c, k: l } = e.inputs, { divisor: u } = e.attrs, d = c.shape[2], f = l.shape[2], p = n(s.bandPart(t([f, f]), -1, 0).cast("bool"), o([f, f]), i([f, f], -Infinity)), m = a(c, l, !1, !0).mul(r(u)), h = p.slice([0, 0], [d, f]).expandDims(0).expandDims(0);
+	return m.add(h);
+}
+e({
+	kernelName: "AttentionMask",
+	backendName: "cpu",
+	kernelFunc: c
+}), e({
+	kernelName: "AttentionMask",
+	backendName: "tensorflow",
+	kernelFunc: c
+});
+//#endregion

package/dist/ops/cpu/fusedSoftmax.d.ts

@@ -0,0 +1,9 @@
+import { SoftmaxAttrs, SoftmaxInputs, TensorInfo } from '@tensorflow/tfjs-core';
+interface FusedSoftmaxAttrs extends SoftmaxAttrs {
+    dropoutRate?: number;
+}
+export declare function softmaxCPU(args: {
+    inputs: SoftmaxInputs;
+    attrs: FusedSoftmaxAttrs;
+}): TensorInfo;
+export {};

package/dist/ops/cpu/fusedSoftmax.js

@@ -0,0 +1,22 @@
+import { Ii as e, Y as t } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/fusedSoftmax.ts
+function n(e) {
+	let { inputs: n, attrs: r } = e, { logits: i } = n, { dim: a, dropoutRate: o } = r;
+	if (!i) throw Error("Error in softmax: input logits is null");
+	return o !== void 0 && o > 0 && console.warn("Dropout in fusedSoftmax not implemented for CPU backend, skipping dropout."), t(i, a);
+}
+e({
+	kernelName: "FusedSoftmax",
+	backendName: "cpu",
+	kernelFunc: n
+}), e({
+	kernelName: "FusedSoftmax",
+	backendName: "tensorflow",
+	kernelFunc: n
+}), e({
+	kernelName: "FusedSoftmax",
+	backendName: "webgpu",
+	kernelFunc: n
+});
+//#endregion
+export { n as softmaxCPU };

package/dist/ops/cpu/gatherSub.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/gatherSub.js

@@ -0,0 +1,12 @@
+import { Bt as e, Ii as t, S as n, V as r, dt as i } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/gatherSub.ts
+function a(t) {
+	let { values: a, labels: o, logits: s } = t.inputs, c = o.shape[0];
+	return e(a, n(s, r([i(0, c, 1, "int32"), o], 1)));
+}
+t({
+	kernelName: "EfficientGatherSub",
+	backendName: "cpu",
+	kernelFunc: a
+});
+//#endregion

package/dist/ops/cpu/gelu.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/gelu.js

@@ -0,0 +1,36 @@
+import { Ii as e, di as t } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/gelu.ts
+var n = .7978845608028654, r = .044715;
+function i(e) {
+	let { inputs: i } = e, { x: a } = i, o = a;
+	return t(() => {
+		let e = o.pow(3), t = o.add(e.mul(r)).mul(n).tanh().add(1).mul(.5);
+		return o.mul(t);
+	});
+}
+e({
+	kernelName: "Gelu",
+	backendName: "cpu",
+	kernelFunc: i
+}), e({
+	kernelName: "Gelu",
+	backendName: "tensorflow",
+	kernelFunc: i
+});
+function a(e) {
+	let { dy: i, x: a } = e.inputs;
+	return t(() => {
+		let e = a.square(), t = e.mul(a), o = a.add(t.mul(r)).mul(n).tanh(), s = o.square().neg().add(1), c = e.mul(3 * r).add(1), l = o.add(1).mul(.5), u = a.mul(s).mul(n).mul(c).mul(.5), d = l.add(u);
+		return i.mul(d);
+	});
+}
+e({
+	kernelName: "GeluGrad",
+	backendName: "cpu",
+	kernelFunc: a
+}), e({
+	kernelName: "GeluGrad",
+	backendName: "tensorflow",
+	kernelFunc: a
+});
+//#endregion

package/dist/ops/cpu/matMul16.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/matMul16.js

@@ -0,0 +1,14 @@
+import { Ii as e, gr as t } from "../../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+//#region lib/ops/cpu/matMul16.ts
+function r(e) {
+	let { A: r, B: i } = e.inputs, { transposeA: a, transposeB: o } = e.attrs, s = !n(r), c = !n(i);
+	if (s && c) return t(r, i, a, o);
+	throw Error("MatMul16 CPU kernel only supports packed tensors currently.");
+}
+e({
+	kernelName: "MatMul16",
+	backendName: "cpu",
+	kernelFunc: r
+});
+//#endregion

package/dist/ops/cpu/matMulGelu.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/matMulGelu.js

@@ -0,0 +1,41 @@
+import { Ii as e, di as t, gr as n } from "../../dist-BewPQWjc.js";
+import { n as r, t as i } from "../../gelu-Bf1HW1RY.js";
+//#region lib/ops/cpu/matMulGelu.ts
+function a(e) {
+	let { inputs: i } = e, { x: a, kernel: o } = i, s = a, c = o;
+	return t(() => r(n(s, c)));
+}
+e({
+	kernelName: "MatMulGelu",
+	backendName: "cpu",
+	kernelFunc: a
+}), e({
+	kernelName: "MatMulGelu",
+	backendName: "tensorflow",
+	kernelFunc: a
+}), e({
+	kernelName: "MatMulGelu",
+	backendName: "webgpu",
+	kernelFunc: a
+});
+function o(e) {
+	let { dy: r, x: a, kernel: o } = e.inputs;
+	return t(() => {
+		let e = i(r, n(a, o));
+		return [e.matMul(o.transpose()), a.transpose().matMul(e)];
+	});
+}
+e({
+	kernelName: "MatMulGeluGrad",
+	backendName: "cpu",
+	kernelFunc: o
+}), e({
+	kernelName: "MatMulGeluGrad",
+	backendName: "tensorflow",
+	kernelFunc: o
+}), e({
+	kernelName: "MatMulGeluGrad",
+	backendName: "webgpu",
+	kernelFunc: o
+});
+//#endregion

package/dist/ops/cpu/matMulMul.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/matMulMul.js

@@ -0,0 +1,20 @@
+import { Ii as e, di as t } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/matMulMul.ts
+function n(e) {
+	let { inputs: n, attrs: r } = e, { transposeA: i, transposeB: a } = r, { x: o, kernel: s, y: c } = n, l = o, u = s, d = c;
+	return t(() => l.matMul(u, i, a).mul(d));
+}
+e({
+	kernelName: "MatMulMul",
+	backendName: "cpu",
+	kernelFunc: n
+}), e({
+	kernelName: "MatMulMul",
+	backendName: "tensorflow",
+	kernelFunc: n
+}), e({
+	kernelName: "MatMulMul",
+	backendName: "webgpu",
+	kernelFunc: n
+});
+//#endregion

package/dist/ops/cpu/mulDropout.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/mulDropout.js

@@ -0,0 +1,20 @@
+import { Ii as e, Wr as t } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/mulDropout.ts
+function n(e) {
+	let { inputs: n } = e, { a: r, b: i } = n;
+	return console.warn("Using fallback mulDrop implementation without dropout."), t(r, i);
+}
+e({
+	kernelName: "MulDropout",
+	backendName: "cpu",
+	kernelFunc: n
+}), e({
+	kernelName: "MulDropout",
+	backendName: "tensorflow",
+	kernelFunc: n
+}), e({
+	kernelName: "MulDropout",
+	backendName: "webgpu",
+	kernelFunc: n
+});
+//#endregion

package/dist/ops/cpu/normRMS.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/normRMS.js

@@ -0,0 +1,35 @@
+import { Ii as e, di as t } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/normRMS.ts
+function n(e) {
+	let { inputs: n } = e, { x: r, gamma: i } = n, a = r, o = i;
+	return t(() => {
+		let e = a.square().mean(-1, !0).add(1e-8).rsqrt();
+		return a.mul(e).mul(o);
+	});
+}
+e({
+	kernelName: "RMSNorm",
+	backendName: "cpu",
+	kernelFunc: n
+}), e({
+	kernelName: "RMSNorm",
+	backendName: "tensorflow",
+	kernelFunc: n
+});
+function r(e) {
+	let { dy: n, x: r, gamma: i } = e.inputs;
+	return t(() => {
+		let e = r.shape[r.shape.length - 1], t = r.square().mean(-1, !0), a = t.add(1e-8).rsqrt(), o = r.mul(a), s = n.mul(o).sum([0, 1]), c = n.mul(i), l = c.mul(r).sum(-1, !0).div(e);
+		return [c.mul(a).sub(r.mul(l).mul(a).div(t.add(1e-8))), s];
+	});
+}
+e({
+	kernelName: "RMSNormGrad",
+	backendName: "cpu",
+	kernelFunc: r
+}), e({
+	kernelName: "RMSNormGrad",
+	backendName: "tensorflow",
+	kernelFunc: r
+});
+//#endregion

package/dist/ops/cpu/qkv.d.ts

@@ -0,0 +1,5 @@
+import { NamedAttrMap, NamedTensorInfoMap, TensorInfo } from '@tensorflow/tfjs-core';
+export declare function qkvCPU(args: {
+    inputs: NamedTensorInfoMap;
+    attrs?: NamedAttrMap;
+}): TensorInfo[];

package/dist/ops/cpu/qkv.js

@@ -0,0 +1,73 @@
+import { G as e, Ii as t, br as n } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/qkv.ts
+function r(t) {
+	let { x: r, kernel: i } = t.inputs, { heads: a, packed: o } = t.attrs;
+	if (o) throw Error("QKV CPU implementation does not support packed tensors.");
+	let [s, c, l] = r.shape, u = n(r, [s * c, l]), d = u.dot(i);
+	u.dispose();
+	let f = n(d, [
+		s,
+		c,
+		3 * l
+	]);
+	d.dispose();
+	let [p, m, h] = e(f, 3, -1);
+	f.dispose();
+	let g = l / a, _ = n(p, [
+		s,
+		c,
+		a,
+		g
+	]);
+	p.dispose();
+	let v = _.transpose([
+		0,
+		2,
+		1,
+		3
+	]);
+	_.dispose();
+	let y = n(m, [
+		s,
+		c,
+		a,
+		g
+	]);
+	m.dispose();
+	let b = y.transpose([
+		0,
+		2,
+		1,
+		3
+	]);
+	y.dispose();
+	let x = n(h, [
+		s,
+		c,
+		a,
+		g
+	]);
+	h.dispose();
+	let S = x.transpose([
+		0,
+		2,
+		1,
+		3
+	]);
+	return x.dispose(), [
+		v,
+		b,
+		S
+	];
+}
+t({
+	kernelName: "QKV",
+	backendName: "cpu",
+	kernelFunc: r
+}), t({
+	kernelName: "QKV",
+	backendName: "tensorflow",
+	kernelFunc: r
+});
+//#endregion
+export { r as qkvCPU };

package/dist/ops/cpu/rope.d.ts

@@ -0,0 +1,6 @@
+import { NamedAttrMap, NamedTensorInfoMap, Tensor, TensorInfo } from '@tensorflow/tfjs-core';
+export declare function applyRoPE(sinCache: Tensor, cosCache: Tensor, rotaryDim: number, q: Tensor, pastLen: number): Tensor;
+export declare function ropeCPU(args: {
+    inputs: NamedTensorInfoMap;
+    attrs?: NamedAttrMap;
+}): TensorInfo;

package/dist/ops/cpu/rope.js

@@ -0,0 +1,81 @@
+import { Ii as e, V as t, _r as n, an as r, dt as i } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/rope.ts
+function a(e, a, o, s, c) {
+	let l = s.shape[3], u = o;
+	if (u > l) return s;
+	let d = s.shape[2], f = u / 2, p = a.slice([
+		c,
+		0,
+		0
+	], [
+		d,
+		f,
+		1
+	]).reshape([
+		1,
+		1,
+		d,
+		f
+	]), m = e.slice([
+		c,
+		0,
+		0
+	], [
+		d,
+		f,
+		1
+	]).reshape([
+		1,
+		1,
+		d,
+		f
+	]), h = s.shape[0], g = s.shape[1], _ = i(0, u, 2, "int32"), v = i(1, u, 2, "int32"), y = ((e) => {
+		let i = e.slice([
+			0,
+			0,
+			0,
+			0
+		], [
+			h,
+			g,
+			d,
+			u
+		]), a = u < l ? e.slice([
+			0,
+			0,
+			0,
+			u
+		], [
+			h,
+			g,
+			d,
+			l - u
+		]) : null, o = r(i, _, 3), s = r(i, v, 3), c = o.mul(p), f = s.mul(m), y = c.sub(f), b = s.mul(p), x = o.mul(m), S = b.add(x);
+		o.dispose(), s.dispose(), p.dispose(), m.dispose(), c.dispose(), f.dispose(), b.dispose(), x.dispose();
+		let C = t([y, S], -1);
+		y.dispose(), S.dispose();
+		let w = C.reshape([
+			h,
+			g,
+			d,
+			u
+		]);
+		return C.dispose(), a ? n([w, a], 3) : w;
+	})(s);
+	return _.dispose(), v.dispose(), y;
+}
+function o(e) {
+	let { x: t } = e.inputs, { pastLen: n, negSin: r, ropeCache: i } = e.attrs, o = t.shape[3];
+	return a(r ? i.getNegSin() : i.getSin(), i.getCos(), o, t, n);
+}
+e({
+	kernelName: "Rope",
+	backendName: "cpu",
+	kernelFunc: o
+}), e({
+	kernelName: "Rope",
+	backendName: "tensorflow",
+	kernelFunc: o
+});
+//#endregion
+export { a as applyRoPE, o as ropeCPU };

package/dist/ops/cpu/scatterSub.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/cpu/scatterSub.js

@@ -0,0 +1,12 @@
+import { Bt as e, C as t, Ii as n, Ot as r, V as i, Wr as a, dt as o } from "../../dist-BewPQWjc.js";
+//#region lib/ops/cpu/scatterSub.ts
+function s(n) {
+	let { logits: s, labels: c, dy: l } = n.inputs, u = c.shape[0], d = s.shape[1];
+	return a(e(s, t(i([o(0, u, 1, "int32"), c], 1), r([u]), [u, d])), l.reshape([u, 1]));
+}
+n({
+	kernelName: "EfficientScatterSub",
+	backendName: "cpu",
+	kernelFunc: s
+});
+//#endregion

package/dist/ops/dot16.d.ts

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function dot16(a: Tensor, b: Tensor, transposeA?: boolean, transposeB?: boolean): Tensor;

package/dist/ops/dot16.js

@@ -0,0 +1,29 @@
+import { o as e } from "../tfjs_backend-BjuQ5FqB.js";
+import { isPackedTensor as t } from "../utilities/packed.js";
+import { transpose16 as n } from "./transpose16.js";
+import { reshape16 as r } from "./reshape16.js";
+import { t as i } from "../matMul16-BNfZSnNM.js";
+//#region lib/ops/dot16.ts
+function a(a, o, s = !1, c = !1) {
+	if (!t(a) && !t(o)) return e(a, o);
+	if (a.rank < 2 || o.rank < 2) throw Error(`dot requires both inputs to be rank >= 2 but got x shape = ${a.shape} and y shape = ${o.shape}`);
+	if (o.rank >= 3 && a.shape.slice(-1)[0] !== o.shape.slice(-2)[0]) throw Error(`If rank y >= 3, then the second last dim of y must equal the last dim of x but got x shape = ${a.shape} and  y shape = ${o.shape}`);
+	if (!t(a) || !t(o)) throw Error("dot16 requires both inputs to be packed Tensors.");
+	if (a.rank === 2 && o.rank === 2) return i(a, o, s, c);
+	{
+		let e = a.shape.slice(), t = e.pop();
+		a = r(a, [-1, t]);
+		let l = o.shape.slice(), u = l.pop(), d = l.pop(), f = [...l, u], p = Array.from({ length: o.rank }, (e, t) => t === 0 ? o.rank - 2 : t <= o.rank - 2 ? t - 1 : t);
+		if (p.every((e, t) => e === t)) o = r(o, [d, -1]);
+		else {
+			let e = n(o, p);
+			o = r(e, [d, -1]), e.dispose();
+		}
+		let m = [...e, ...f], h = i(a, o, s, c);
+		a.dispose(), o.dispose();
+		let g = r(h, m);
+		return h.dispose(), g;
+	}
+}
+//#endregion
+export { a as dot16 };

package/dist/ops/dropout.d.ts

@@ -0,0 +1,2 @@
+import { Tensor, TensorLike } from '@tensorflow/tfjs-core';
+export declare function dropout(x: Tensor | TensorLike, rate: number, noiseShape?: number[], seed?: number | string): Tensor;

package/dist/ops/dropout.js

@@ -0,0 +1,11 @@
+import { Gr as e, Ps as t, ft as n, mi as r, on as i, qr as a, wi as o, x as s } from "../dist-BewPQWjc.js";
+import { mul16 as c } from "./mul16.js";
+//#region lib/ops/dropout.ts
+function l(l, u, d, f) {
+	let p = r(l, "x", "dropout");
+	if (t(u >= 0 && u < 1, () => `rate must be a float in the range [0, 1), but got ${u}.`), u === 0) return l instanceof o ? p.clone() : p;
+	let m = s(p, d), h = 1 - u;
+	return c(p, e(i(a(n(m, 0, 1, "float32", f), h)), h));
+}
+//#endregion
+export { l as dropout };

package/dist/ops/dropout16.d.ts

@@ -0,0 +1,2 @@
+import { Rank, Tensor } from '@tensorflow/tfjs-core';
+export declare function dropout16<R extends Rank = Rank>(x: Tensor<R>, dropout: number, seed?: number): Tensor<R>;

package/dist/ops/dropout16.js

@@ -0,0 +1,22 @@
+import { Fi as e, ii as t } from "../dist-BewPQWjc.js";
+import { isPackedTensor as n } from "../utilities/packed.js";
+import "./webgl/dropout16.js";
+//#region lib/ops/grads/dropout16.ts
+e({
+	kernelName: "Dropout16",
+	inputsToSave: [],
+	gradFunc: (e, t, n) => {
+		let { dropout: i, seed: a } = n;
+		return { x: () => r(e, i, a) };
+	}
+});
+//#endregion
+//#region lib/ops/dropout16.ts
+function r(e, r, i) {
+	return n(e) ? t().runKernel("Dropout16", { x: e }, {
+		dropout: r,
+		seed: i ?? Math.random()
+	}) : e;
+}
+//#endregion
+export { r as dropout16 };

package/dist/ops/gatherSub.d.ts

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function gatherSub(values: Tensor, labels: Tensor, logits: Tensor): Tensor;

package/dist/ops/gatherSub.js

@@ -0,0 +1,13 @@
+import { ii as e } from "../dist-BewPQWjc.js";
+import "./cpu/gatherSub.js";
+import "./webgl/gatherSub.js";
+//#region lib/ops/gatherSub.ts
+function t(t, n, r) {
+	return e().runKernel("EfficientGatherSub", {
+		logits: r,
+		labels: n,
+		values: t
+	}, {});
+}
+//#endregion
+export { t as gatherSub };

package/dist/ops/gelu.d.ts

@@ -0,0 +1,3 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function gelu(x: Tensor): Tensor;
+export declare function dGelu(dy: Tensor, x: Tensor): Tensor;

package/dist/ops/gelu.js

@@ -0,0 +1,2 @@
+import { n as e, t } from "../gelu-Bf1HW1RY.js";
+export { t as dGelu, e as gelu };

package/dist/ops/globalNorm.d.ts

@@ -0,0 +1,2 @@
+import { Scalar, Tensor } from '@tensorflow/tfjs-core';
+export declare function clipScale(grads: Tensor[], invLossScaling: number, clipNorm?: number): Scalar;

package/dist/ops/globalNorm.js

@@ -0,0 +1,19 @@
+import { di as e, ii as t, kt as n } from "../dist-BewPQWjc.js";
+//#region lib/ops/globalNorm.ts
+function r(r, i, a = 1) {
+	return e(() => {
+		let e = n([r.length], "int32");
+		return r.map((n, r) => t().runKernel("Norm2", {
+			x: n,
+			output: e
+		}, {
+			invLossScaling: i,
+			index: r
+		})), t().runKernel("ClipScale", { x: e }, {
+			invLossScaling: i,
+			clipNorm: a
+		});
+	});
+}
+//#endregion
+export { r as clipScale };

package/dist/ops/grads/add16.d.ts

@@ -0,0 +1 @@
+export {};