@genai-fi/nanogpt 0.2.11 → 0.3.0

package/dist/moments-DfcpfwKi.js

@@ -0,0 +1,132 @@
+import { o as i, h as m, E as y, a2 as A, a0 as E, a3 as o, a4 as h, J as b, f as S, a5 as v, a6 as _, s as k, U as q } from "./index-pWA4_lUh.js";
+import { e as w, s as u, m as l } from "./sum-C7Mgy9Bw.js";
+import { r as c } from "./reshape-C8CR_Bad.js";
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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 K(r, e = null, n = !1) {
+  const t = { x: m(r, "x", "min") }, a = { axis: e, keepDims: n };
+  return y.runKernel(A, t, a);
+}
+const d = /* @__PURE__ */ i({ min_: K });
+/**
+ * @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 M(r, e = "euclidean", n = null, s = !1) {
+  r = m(r, "x", "norm");
+  const t = $(r, e, n);
+  let a = t.shape;
+  if (s) {
+    const f = E(n, r.shape);
+    a = w(t.shape, f);
+  }
+  return c(t, a);
+}
+function $(r, e, n = null) {
+  if (r.rank === 0)
+    return o(r);
+  if (r.rank !== 1 && n === null)
+    return $(c(r, [-1]), e, n);
+  if (r.rank === 1 || typeof n == "number" || Array.isArray(n) && n.length === 1) {
+    if (e === 1)
+      return u(o(r), n);
+    if (e === 1 / 0)
+      return l(o(r), n);
+    if (e === -1 / 0)
+      return d(o(r), n);
+    if (e === "euclidean" || e === 2)
+      return h(u(b(o(r), S(2, "int32")), n));
+    throw new Error(`Error in norm: invalid ord value: ${e}`);
+  }
+  if (Array.isArray(n) && n.length === 2) {
+    if (e === 1)
+      return l(u(o(r), n[0]), n[1] - 1);
+    if (e === 1 / 0)
+      return l(u(o(r), n[1]), n[0]);
+    if (e === -1 / 0)
+      return d(u(o(r), n[1]), n[0]);
+    if (e === "fro" || e === "euclidean")
+      return h(u(v(r), n));
+    throw new Error(`Error in norm: invalid ord value: ${e}`);
+  }
+  throw new Error(`Error in norm: invalid axis: ${n}`);
+}
+const J = /* @__PURE__ */ i({ norm_: M });
+/**
+ * @license
+ * Copyright 2020 Google Inc. 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 T(r, e = null, n = !1) {
+  const t = { x: m(r, "x", "mean") }, a = { axis: e, keepDims: n };
+  return y.runKernel(_, t, a);
+}
+const p = /* @__PURE__ */ i({ mean_: T });
+/**
+ * @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 g(r, e = null, n = !1) {
+  r = m(r, "x", "moments");
+  const s = E(e, r.shape), t = p(r, s, n);
+  let a = t.shape;
+  n || (a = w(t.shape, s));
+  const f = v(k(q(r, "float32"), c(t, a))), I = p(f, s, n);
+  return { mean: t, variance: I };
+}
+const P = /* @__PURE__ */ i({ moments_: g });
+export {
+  P as a,
+  d as b,
+  p as m,
+  J as n
+};

package/dist/ones-Cog-G2ag.js

@@ -0,0 +1,29 @@
+import { k as n, l as t, n as m, E as i } from "./index-pWA4_lUh.js";
+import { z as l, c } from "./zeros-CCy9C3uU.js";
+/**
+ * @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 f(o, r = "float32") {
+  if (n(o), r === "complex64") {
+    const s = f(o, "float32"), a = l(o, "float32");
+    return c(s, a);
+  }
+  const e = t(m(o), r);
+  return i.makeTensor(e, o, r);
+}
+export {
+  f as o
+};

package/dist/ops/appendCache.d.ts

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function appendCache(cache: Tensor, item: Tensor, maxSize: number): Tensor;

package/dist/ops/appendCache.js

@@ -0,0 +1,9 @@
+import { e as p } from "../index-pWA4_lUh.js";
+import "./cpu/appendCache.js";
+import "./webgl/appendCache.js";
+function a(e, n, r) {
+  return p().runKernel("AppendCache", { cache: e, item: n }, { maxSize: r });
+}
+export {
+  a as appendCache
+};

package/dist/ops/attentionMask.d.ts

@@ -1,2 +1,2 @@
-import { Tensor } from '@tensorflow/tfjs';
+import { Tensor } from '@tensorflow/tfjs-core';
 export declare function attentionMask(q: Tensor, k: Tensor, mask: Tensor, divisor: number): Tensor;

package/dist/ops/attentionMask.js

@@ -1,88 +1,10 @@
-import { engine as k } from "@tensorflow/tfjs";
-import { r as m, c as d, s as p } from "../index-YPKosni4.js";
-import { m as f } from "../mat_mul-Bu7bhLms.js";
-class h {
-  variableNames = ["q", "k", "mask"];
-  outputShape;
-  userCode;
-  // enableShapeUniforms = true;
-  customUniforms = [{ name: "divisor", type: "float" }];
-  constructor(e, n, s, a) {
-    this.outputShape = [e, n, s, s], this.userCode = `
-        void main() {
-            ivec4 coords = getOutputCoords(); // [batch, nh, t1, t2]
-            int b = coords.x;
-            int h = coords.y;
-            int t1 = coords.z;
-            int t2 = coords.w;
-
-            float sum = 0.0;
-            for (int i = 0; i < ${a}; ++i) {
-                float qv = getQ(b, h, t1, i);
-                float kv = getK(b, h, t2, i); // k is transposed on last two dims
-                sum += qv * kv;
-            }
-
-            // Scale by divisor
-            float scaled = sum * divisor;
-
-            // Add mask
-            float maskVal = getMask(t1, t2); // mask is [T,T]
-
-            setOutput(scaled + maskVal);
-        }
-        `;
-  }
+import { e as o } from "../index-pWA4_lUh.js";
+import "./cpu/attentionMask.js";
+import "./webgl/attentionMask.js";
+import "./grads/attentionMask.js";
+function s(t, n, r, e) {
+  return o().runKernel("AttentionMask", { q: t, k: n, mask: r }, { divisor: e });
 }
-function v(t) {
-  const { q: e, k: n, mask: s } = t.inputs, { divisor: a } = t.attrs, o = t.backend, r = e.shape[0], i = e.shape[2], c = e.shape[1], u = new h(r, c, i, e.shape[3]);
-  return o.runWebGLProgram(u, [e, n, s], "float32", [[a]]);
-}
-const b = {
-  kernelName: "AttentionMask",
-  backendName: "webgl",
-  kernelFunc: v
-};
-m(b);
-function l(t) {
-  const { q: e, k: n, mask: s } = t.inputs, { divisor: a } = t.attrs, o = e.shape[2], i = f(e, n, !1, !0).mul(p(a)), c = s.slice([0, 0], [o, o]).expandDims(0).expandDims(0);
-  return i.add(c);
-}
-const M = {
-  kernelName: "AttentionMask",
-  backendName: "cpu",
-  kernelFunc: l
-};
-m(M);
-const g = {
-  kernelName: "AttentionMask",
-  backendName: "tensorflow",
-  kernelFunc: l
-};
-m(g);
-function N(t, e, n, s) {
-  return k().runKernel("AttentionMask", { q: t, k: e, mask: n }, { divisor: s });
-}
-const A = {
-  kernelName: "AttentionMask",
-  inputsToSave: ["q", "k"],
-  outputsToSave: [],
-  gradFunc: (t, e, n) => {
-    if (Array.isArray(t))
-      throw new Error("Expected dy to be a single Tensor");
-    const [s, a] = e, { divisor: o } = n;
-    return {
-      q: () => t.matMul(a).mul(o),
-      k: () => s.transpose([0, 1, 3, 2]).matMul(t).mul(o).transpose([0, 1, 3, 2]),
-      mask: () => t,
-      divisor: () => {
-        const r = s.matMul(a, !1, !0);
-        return t.mul(r).sum();
-      }
-    };
-  }
-};
-d(A);
 export {
-  N as attentionMask
+  s as attentionMask
 };

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

@@ -0,0 +1,2 @@
+import { Tensor } from '@tensorflow/tfjs-core';
+export declare function appendCache(cache: Tensor, item: Tensor, maxSize: number): Tensor;

package/dist/ops/cpu/appendCache.js

@@ -0,0 +1,28 @@
+import { r as a, e as m } from "../../index-pWA4_lUh.js";
+import { c as d } from "../../concat-BIZS_td9.js";
+function r(n) {
+  const { cache: c, item: t } = n.inputs, { maxSize: o } = n.attrs, e = d([c, t], 2), s = e.shape[2];
+  if (s > o) {
+    const p = s - o, i = e.shape[0], l = e.shape[1], h = e.shape[3], u = e.slice([0, 0, p, 0], [i, l, o, h]);
+    return e.dispose(), u;
+  }
+  return e;
+}
+const f = {
+  kernelName: "AppendCache",
+  backendName: "cpu",
+  kernelFunc: r
+};
+a(f);
+const C = {
+  kernelName: "AppendCache",
+  backendName: "tensorflow",
+  kernelFunc: r
+};
+a(C);
+function N(n, c, t) {
+  return m().runKernel("AppendCache", { cache: n, item: c }, { maxSize: t });
+}
+export {
+  N as appendCache
+};

package/dist/ops/cpu/attentionMask.js

@@ -0,0 +1,18 @@
+import { r as s, f as i } from "../../index-pWA4_lUh.js";
+import { m as k } from "../../mat_mul-D7_a4KJn.js";
+function a(t) {
+  const { q: e, k: o, mask: r } = t.inputs, { divisor: c } = t.attrs, n = e.shape[2], m = k(e, o, !1, !0).mul(i(c)), l = r.slice([0, 0], [n, n]).expandDims(0).expandDims(0);
+  return m.add(l);
+}
+const d = {
+  kernelName: "AttentionMask",
+  backendName: "cpu",
+  kernelFunc: a
+};
+s(d);
+const u = {
+  kernelName: "AttentionMask",
+  backendName: "tensorflow",
+  kernelFunc: a
+};
+s(u);

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

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

package/dist/ops/cpu/gatherSub.js

@@ -0,0 +1,34 @@
+import { o as u, h as c, E as g, H as h, r as m, s as p } from "../../index-pWA4_lUh.js";
+import { r as l } from "../../range-CcDl05lo.js";
+import { s as N } from "../../stack-BMm-efee.js";
+/**
+ * @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 f(e, s) {
+  const n = c(s, "indices", "gatherND", "int32"), t = { params: c(e, "x", "gatherND", "string_or_numeric"), indices: n };
+  return g.runKernel(h, t);
+}
+const b = /* @__PURE__ */ u({ gatherND_: f });
+function d(e) {
+  const { values: s, labels: n, logits: r } = e.inputs, t = n.shape[0], o = l(0, t, 1, "int32"), a = N([o, n], 1), i = b(r, a);
+  return p(s, i);
+}
+const k = {
+  kernelName: "EfficientGatherSub",
+  backendName: "cpu",
+  kernelFunc: d
+};
+m(k);

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,38 @@
+import { r as q } from "../../index-pWA4_lUh.js";
+import { r as o } from "../../reshape-C8CR_Bad.js";
+import { s as x } from "../../split-DZbvruEP.js";
+function v(p) {
+  const { x: c, kernel: K } = p.inputs, { heads: n } = p.attrs, [s, e, t] = c.shape, a = o(c, [s * e, t]), i = a.dot(K);
+  a.dispose();
+  const d = o(i, [s, e, 3 * t]);
+  i.dispose();
+  const [k, l, m] = x(d, 3, -1);
+  d.dispose();
+  const r = t / n, f = o(k, [s, e, n, r]);
+  k.dispose();
+  const C = f.transpose([0, 2, 1, 3]);
+  f.dispose();
+  const h = o(l, [s, e, n, r]);
+  l.dispose();
+  const N = h.transpose([0, 2, 1, 3]);
+  h.dispose();
+  const u = o(m, [s, e, n, r]);
+  m.dispose();
+  const T = u.transpose([0, 2, 1, 3]);
+  return u.dispose(), [C, N, T];
+}
+const F = {
+  kernelName: "QKV",
+  backendName: "cpu",
+  kernelFunc: v
+};
+q(F);
+const R = {
+  kernelName: "QKV",
+  backendName: "tensorflow",
+  kernelFunc: v
+};
+q(R);
+export {
+  v 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,38 @@
+import { r as S } from "../../index-pWA4_lUh.js";
+import { r as F } from "../../range-CcDl05lo.js";
+import { g as I } from "../../gather-BPGW8RsB.js";
+import { s as E } from "../../stack-BMm-efee.js";
+import { c as T } from "../../concat-BIZS_td9.js";
+function U(t, c, p, o, r) {
+  const n = o.shape[3], s = p;
+  if (s > n) return o;
+  const e = o.shape[2], i = s / 2, a = c.slice([r, 0, 0], [e, i, 1]).reshape([1, 1, e, i]), d = t.slice([r, 0, 0], [e, i, 1]).reshape([1, 1, e, i]), l = o.shape[0], m = o.shape[1], h = F(0, s, 2, "int32"), k = F(1, s, 2, "int32"), D = ((R) => {
+    const g = R.slice([0, 0, 0, 0], [l, m, e, s]), v = s < n ? R.slice([0, 0, 0, s], [l, m, e, n - s]) : null, u = I(g, h, 3), f = I(g, k, 3), C = u.mul(a), N = f.mul(d), P = C.sub(N), b = f.mul(a), x = u.mul(d), K = b.add(x);
+    u.dispose(), f.dispose(), a.dispose(), d.dispose(), C.dispose(), N.dispose(), b.dispose(), x.dispose();
+    const w = E([P, K], -1);
+    P.dispose(), K.dispose();
+    const y = w.reshape([l, m, e, s]);
+    return w.dispose(), v ? T([y, v], 3) : y;
+  })(o);
+  return h.dispose(), k.dispose(), D;
+}
+function B(t) {
+  const { x: c, sin: p, cos: o } = t.inputs, { pastLen: r } = t.attrs, n = c.shape[3];
+  return U(p, o, n, c, r);
+}
+const j = {
+  kernelName: "Rope",
+  backendName: "cpu",
+  kernelFunc: B
+};
+S(j);
+const z = {
+  kernelName: "Rope",
+  backendName: "tensorflow",
+  kernelFunc: B
+};
+S(z);
+export {
+  U as applyRoPE,
+  B as ropeCPU
+};

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

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

package/dist/ops/cpu/scatterSub.js

@@ -0,0 +1,70 @@
+import { o as l, k, h, E as g, Q as w, r as $, s as d, b as m } from "../../index-pWA4_lUh.js";
+import { r as b } from "../../range-CcDl05lo.js";
+import { s as E } from "../../stack-BMm-efee.js";
+import { o as D } from "../../ones-Cog-G2ag.js";
+function N(a, r, t) {
+  const s = r.rank > 1 ? r.shape[r.rank - 1] : 1, e = r.rank > 1 ? r.rank - 1 : 1, o = `Must have updates.shape = indices.shape[:batchDim] + shape[sliceDim:], got updates.shape: ${t.shape}, indices.shape: ${r.shape}, shape: ${a}, sliceDim: ${s}, and batchDim: ${e}.`;
+  if (t.rank < e)
+    throw new Error(o + ` update.rank < ${e}. `);
+  if (a.length < s + (t.rank - e))
+    throw new Error(o + ` Output shape length < ${s + (t.rank - e)}`);
+  if (t.rank !== e + a.length - s)
+    throw new Error(o + ` update.rank != ${e + a.length - s}`);
+  for (let n = 0; n < e; ++n)
+    if (t.shape[n] !== r.shape[n])
+      throw new Error(o + ` updates.shape[${n}] (${t.shape[n]}) != indices.shape[${n}] (${r.shape[n]}).`);
+  for (let n = 0; n < t.rank - e; ++n)
+    if (t.shape[n + e] !== a[n + s])
+      throw new Error(o + ` updates.shape[${n + e}] (${t.shape[n + e]}) != shape[${n + e}] (${a[n + e]})`);
+}
+function S(a, r, t) {
+  if (r.rank < 1)
+    throw new Error(`tf.scatterND() expects the indices to be rank 1 or higher, but the rank was ${r.rank}.`);
+  if (a.rank < 1)
+    throw new Error(`tf.scatterND() expects the updates to be rank 1 or higher, but the rank was ${a.rank}.`);
+  if (r.dtype !== "int32")
+    throw new Error(`The dtype of 'indices' should be int32, but got dtype: ${r.dtype}`);
+  if (t.length < 1)
+    throw new Error(`Output rank must be greater or equal to 1, but got shape: ${t}`);
+  if (t.length === 0) {
+    if (r.size === 0)
+      throw new Error(`Indices specified for empty output. indices shape: ${r.shape}`);
+    if (a.size === 0)
+      throw new Error(`Updates specified for empty output. updates shape: ${a.shape}`);
+  }
+  N(t, r, 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 y(a, r, t) {
+  k(t);
+  const s = h(a, "indices", "scatterND", "int32"), e = h(r, "updates", "scatterND");
+  S(e, s, t);
+  const o = { indices: s, updates: e }, n = { shape: t };
+  return g.runKernel(w, o, n);
+}
+const v = /* @__PURE__ */ l({ scatterND_: y });
+function I(a) {
+  const { logits: r, labels: t, dy: s } = a.inputs, e = t.shape[0], o = r.shape[1], n = b(0, e, 1, "int32"), i = E([n, t], 1), c = D([e]), p = v(i, c, [e, o]), f = d(r, p), u = s.reshape([e, 1]);
+  return m(f, u);
+}
+const T = {
+  kernelName: "EfficientScatterSub",
+  backendName: "cpu",
+  kernelFunc: I
+};
+$(T);

package/dist/ops/gatherSub.d.ts

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

package/dist/ops/gatherSub.js

@@ -1,66 +1,9 @@
-import { engine as l } from "@tensorflow/tfjs";
-import { o as g, d as i, E as b, G as d, r as c, b as h } from "../index-YPKosni4.js";
-import { r as p, s as f } from "../stack-BtKpB0Ry.js";
-/**
- * @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 m(e, t) {
-  const n = i(t, "indices", "gatherND", "int32"), r = { params: i(e, "x", "gatherND", "string_or_numeric"), indices: n };
-  return b.runKernel(d, r);
-}
-const N = /* @__PURE__ */ g({ gatherND_: m });
-class S {
-  variableNames = ["labels", "logits", "values"];
-  outputShape;
-  userCode;
-  constructor(t) {
-    this.outputShape = [t], this.userCode = `
-      void main() {
-        int coords = getOutputCoords();
-        int index = int(getLabelsAtOutCoords());
-        float val = getValuesAtOutCoords();
-        float logit = getLogits(coords, index);
-        setOutput(val - logit);
-      }
-    `;
-  }
-}
-function k(e) {
-  const { logits: t, labels: n, values: s } = e.inputs, r = e.backend, o = n.shape[0], a = new S(o);
-  return r.runWebGLProgram(a, [n, t, s], "float32");
-}
-const G = {
-  kernelName: "EfficientGatherSub",
-  backendName: "webgl",
-  kernelFunc: k
-};
-c(G);
-function v(e) {
-  const { values: t, labels: n, logits: s } = e.inputs, r = n.shape[0], o = p(0, r, 1, "int32"), a = f([o, n], 1), u = N(s, a);
-  return h(t, u);
-}
-const C = {
-  kernelName: "EfficientGatherSub",
-  backendName: "cpu",
-  kernelFunc: v
-};
-c(C);
-function K(e, t, n) {
-  return l().runKernel("EfficientGatherSub", { logits: n, labels: t, values: e }, {});
+import { e as n } from "../index-pWA4_lUh.js";
+import "./cpu/gatherSub.js";
+import "./webgl/gatherSub.js";
+function f(r, e, t) {
+  return n().runKernel("EfficientGatherSub", { logits: t, labels: e, values: r }, {});
 }
 export {
-  K as gatherSub
+  f as gatherSub
 };

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

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

package/dist/ops/grads/attentionMask.js

@@ -0,0 +1,21 @@
+import { g as i } from "../../index-pWA4_lUh.js";
+const u = {
+  kernelName: "AttentionMask",
+  inputsToSave: ["q", "k"],
+  outputsToSave: [],
+  gradFunc: (t, a, n) => {
+    if (Array.isArray(t))
+      throw new Error("Expected dy to be a single Tensor");
+    const [r, e] = a, { divisor: s } = n;
+    return {
+      q: () => t.matMul(e).mul(s),
+      k: () => r.transpose([0, 1, 3, 2]).matMul(t).mul(s).transpose([0, 1, 3, 2]),
+      mask: () => t,
+      divisor: () => {
+        const o = r.matMul(e, !1, !0);
+        return t.mul(o).sum();
+      }
+    };
+  }
+};
+i(u);

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

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

package/dist/ops/grads/qkv.js

@@ -0,0 +1,20 @@
+import { g as v } from "../../index-pWA4_lUh.js";
+const g = {
+  kernelName: "QKV",
+  inputsToSave: ["x", "kernel"],
+  outputsToSave: [],
+  gradFunc: (k, m) => {
+    const [x, K, f] = k, [c, r] = m, [t, s, e] = c.shape, d = x.transpose([0, 2, 1, 3]).reshape([t * s, e]), l = K.transpose([0, 2, 1, 3]).reshape([t * s, e]), u = f.transpose([0, 2, 1, 3]).reshape([t * s, e]), i = r.slice([0, 0], [e, e]), h = r.slice([0, e], [e, e]), M = r.slice([0, 2 * e], [e, e]);
+    return {
+      x: () => {
+        const n = d.matMul(i, !1, !0), a = l.matMul(h, !1, !0), o = u.matMul(M, !1, !0);
+        return n.add(a).add(o).reshape([t, s, e]);
+      },
+      kernel: () => {
+        const n = c.reshape([t * s, e]), a = n.matMul(d, !0, !1), o = n.matMul(l, !0, !1), p = n.matMul(u, !0, !1);
+        return a.concat(o, 1).concat(p, 1);
+      }
+    };
+  }
+};
+v(g);

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

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