@genai-fi/nanogpt 0.2.12 → 0.3.0

package/dist/ops/grads/rope.js

@@ -0,0 +1,14 @@
+import { g as a, e as i } from "../../index-pWA4_lUh.js";
+function p(n, e, s, o) {
+  return i().runKernel("Rope", { x: n, sin: e, cos: s }, { pastLen: o });
+}
+const c = {
+  kernelName: "Rope",
+  inputsToSave: ["sin", "cos"],
+  outputsToSave: [],
+  gradFunc: (n, e) => {
+    const [s, o] = e, t = s.neg(), r = p(n, t, o, 0);
+    return t.dispose(), { x: () => r };
+  }
+};
+a(c);

package/dist/ops/node/sparseCrossEntropy.js

@@ -1,4 +1,4 @@
-import { r as o } from "../../index-YPKosni4.js";
+import { r as o } from "../../index-pWA4_lUh.js";
 function r(e) {
   const { logits: t, labels: n } = e.inputs;
   return e.backend.executeMultipleOutputs("SparseSoftmaxCrossEntropyWithLogits", [], [t, n], 2);

package/dist/ops/qkv.d.ts

@@ -1,7 +1,2 @@
-import { Tensor } from '@tensorflow/tfjs';
-import { TensorInfo, NamedTensorInfoMap, NamedAttrMap } from '@tensorflow/tfjs-core';
-export declare function qkvCPU(args: {
-    inputs: NamedTensorInfoMap;
-    attrs?: NamedAttrMap;
-}): TensorInfo[];
+import { Tensor } from '@tensorflow/tfjs-core';
 export declare function qkv(x: Tensor, kernel: Tensor, heads: number): Tensor[];

package/dist/ops/qkv.js

@@ -1,127 +1,10 @@
-import { engine as C } from "@tensorflow/tfjs";
-import { o as N, d as V, E as Q, i as T, r as b, c as G } from "../index-YPKosni4.js";
-import { r as x } from "../reshape-DmnmKT6r.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.
- * =============================================================================
- */
-function M(r, e, o = 0) {
-  const t = { x: V(r, "x", "split") }, s = { numOrSizeSplits: e, axis: o };
-  return Q.runKernel(T, t, s);
+import { e as o } from "../index-pWA4_lUh.js";
+import "./cpu/qkv.js";
+import "./webgl/qkv.js";
+import "./grads/qkv.js";
+function u(r, e, n) {
+  return o().runKernel("QKV", { x: r, kernel: e }, { heads: n });
 }
-const P = /* @__PURE__ */ N({ split_: M });
-class S {
-  variableNames = ["x", "kernel"];
-  outputShape;
-  userCode;
-  // enableShapeUniforms = true;
-  customUniforms = [{ name: "mode", type: "int" }];
-  constructor(e, o, a, t) {
-    const s = t / o;
-    this.outputShape = [e, o, a, s], this.userCode = `
-        void main() {
-            ivec4 coords = getOutputCoords(); // [b, h, t, d]
-            int b = coords.x;
-            int h = coords.y;
-            int t = coords.z;
-            int d = coords.w;
-
-            // Compute output channel index in fused kernel
-            int out_offset = mode * ${o} * ${s} + h * ${s} + d;
-
-            float sum = 0.0;
-            for (int c = 0; c < ${t}; ++c) {
-                float xval = getX(b, t, c); // fetch from x
-                float kval = getKernel(c, out_offset); // fetch from kernel
-                sum += xval * kval;
-            }
-
-            setOutput(sum);
-        }
-        `;
-  }
-}
-function $(r) {
-  const { x: e, kernel: o } = r.inputs, { heads: a } = r.attrs, t = r.backend, s = e.shape[0], l = e.shape[1], d = e.shape[2], c = new S(s, a, l, d);
-  return [
-    t.runWebGLProgram(c, [e, o], "float32", [[0]]),
-    t.runWebGLProgram(c, [e, o], "float32", [[1]]),
-    t.runWebGLProgram(c, [e, o], "float32", [[2]])
-  ];
-}
-const w = {
-  kernelName: "QKV",
-  backendName: "webgl",
-  kernelFunc: $
-};
-b(w);
-function q(r) {
-  const { x: e, kernel: o } = r.inputs, { heads: a } = r.attrs, [t, s, l] = e.shape, d = x(e, [t * s, l]), c = d.dot(o);
-  d.dispose();
-  const n = x(c, [t, s, 3 * l]);
-  c.dispose();
-  const [p, m, k] = P(n, 3, -1);
-  n.dispose();
-  const f = l / a, v = x(p, [t, s, a, f]);
-  p.dispose();
-  const g = v.transpose([0, 2, 1, 3]);
-  v.dispose();
-  const u = x(m, [t, s, a, f]);
-  m.dispose();
-  const h = u.transpose([0, 2, 1, 3]);
-  u.dispose();
-  const i = x(k, [t, s, a, f]);
-  k.dispose();
-  const K = i.transpose([0, 2, 1, 3]);
-  return i.dispose(), [g, h, K];
-}
-const F = {
-  kernelName: "QKV",
-  backendName: "cpu",
-  kernelFunc: q
-};
-b(F);
-const L = {
-  kernelName: "QKV",
-  backendName: "tensorflow",
-  kernelFunc: q
-};
-b(L);
-function y(r, e, o) {
-  return C().runKernel("QKV", { x: r, kernel: e }, { heads: o });
-}
-const _ = {
-  kernelName: "QKV",
-  inputsToSave: ["x", "kernel"],
-  outputsToSave: [],
-  gradFunc: (r, e) => {
-    const [o, a, t] = r, [s, l] = e, [d, c, n] = s.shape, p = o.transpose([0, 2, 1, 3]).reshape([d * c, n]), m = a.transpose([0, 2, 1, 3]).reshape([d * c, n]), k = t.transpose([0, 2, 1, 3]).reshape([d * c, n]), f = l.slice([0, 0], [n, n]), v = l.slice([0, n], [n, n]), g = l.slice([0, 2 * n], [n, n]);
-    return {
-      x: () => {
-        const u = p.matMul(f, !1, !0), h = m.matMul(v, !1, !0), i = k.matMul(g, !1, !0);
-        return u.add(h).add(i).reshape([d, c, n]);
-      },
-      kernel: () => {
-        const u = s.reshape([d * c, n]), h = u.matMul(p, !0, !1), i = u.matMul(m, !0, !1), K = u.matMul(k, !0, !1);
-        return h.concat(i, 1).concat(K, 1);
-      }
-    };
-  }
-};
-G(_);
 export {
-  y as qkv,
-  q as qkvCPU
+  u as qkv
 };

package/dist/ops/rope.d.ts

@@ -1,8 +1,3 @@
 import { default as RoPECache } from '../layers/RoPECache';
 import { Tensor } from '@tensorflow/tfjs';
-import { TensorInfo, NamedTensorInfoMap, NamedAttrMap } from '@tensorflow/tfjs-core';
-export declare function ropeCPU(args: {
-    inputs: NamedTensorInfoMap;
-    attrs?: NamedAttrMap;
-}): TensorInfo;
 export declare function rope(x: Tensor, cache: RoPECache, pastLength: number): Tensor;

package/dist/ops/rope.js

@@ -1,154 +1,10 @@
-import { engine as D } from "@tensorflow/tfjs";
-import { o as G, l as F, k as _, n as U, E as K, p as z, d as I, q as O, r as f, c as A } from "../index-YPKosni4.js";
-import { r as T, s as B } from "../stack-BtKpB0Ry.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.
- * =============================================================================
- */
-function W(o, t = 0) {
-  F(o.length >= 1, () => "Pass at least one tensor to concat");
-  const e = _(o, "tensors", "concat", "string_or_numeric");
-  if (e[0].dtype === "complex64" && e.forEach((r) => {
-    if (r.dtype !== "complex64")
-      throw new Error(`Cannot concatenate complex64 tensors with a tensor
-          with dtype ${r.dtype}. `);
-  }), e.length === 1)
-    return U(e[0]);
-  const n = e, s = { axis: t };
-  return K.runKernel(z, n, s);
+import { engine as n } from "@tensorflow/tfjs";
+import "./cpu/rope.js";
+import "./webgl/rope.js";
+import "./grads/rope.js";
+function s(o, e, r) {
+  return e.ensureRopeCache(o.shape[1] + r), n().runKernel("Rope", { x: o, sin: e.getSin(), cos: e.getCos() }, { pastLen: r });
 }
-const j = /* @__PURE__ */ G({ concat_: W });
-/**
- * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- * =============================================================================
- */
-function H(o, t, e = 0, n = 0) {
-  const s = I(o, "x", "gather"), r = I(t, "indices", "gather", "int32"), c = { x: s, indices: r }, a = { axis: e, batchDims: n };
-  return K.runKernel(O, c, a);
-}
-const E = /* @__PURE__ */ G({ gather_: H });
-class J {
-  variableNames = ["x", "sin", "cos"];
-  outputShape;
-  userCode;
-  // enableShapeUniforms = true;
-  customUniforms = [{ name: "pastLen", type: "int" }];
-  constructor(t, e, n, s) {
-    this.outputShape = [t, e, n, s], this.userCode = `
-        void main() {
-            ivec4 coords = getOutputCoords(); // [b, h, t, d]
-            int b = coords.x;
-            int h = coords.y;
-            int t = coords.z;
-            int d = coords.w;
-
-            int rotaryDim = ${s};
-
-            float outVal = 0.0;
-
-            if (d < rotaryDim) {
-                int pairIdx = d / 2;
-                float cos = getCos(t + pastLen, pairIdx, 0);
-                float sin = getSin(t + pastLen, pairIdx, 0);
-
-                if (d % 2 == 0) {
-                    // even index
-                    float even = getX(b, h, t, d);
-                    float odd = getX(b, h, t, d + 1);
-                    outVal = even * cos - odd * sin;
-                } else {
-                    // odd index
-                    float even = getX(b, h, t, d - 1);
-                    float odd = getX(b, h, t, d);
-                    outVal = even * sin + odd * cos;
-                }
-            } else {
-                // pass through for non-rotary dims
-                outVal = getX(b, h, t, d);
-            }
-
-            setOutput(outVal);
-        }
-        `;
-  }
-}
-function M(o) {
-  const { x: t, sin: e, cos: n } = o.inputs, { pastLen: s } = o.attrs, r = o.backend, c = t.shape[0], a = t.shape[1], i = t.shape[2], d = t.shape[3], p = new J(c, a, i, d);
-  return r.runWebGLProgram(p, [t, e, n], "float32", [[s]]);
-}
-const Q = {
-  kernelName: "Rope",
-  backendName: "webgl",
-  kernelFunc: M
-};
-f(Q);
-function V(o, t, e, n, s) {
-  const r = n.shape[3], c = e;
-  if (c > r) return n;
-  const a = n.shape[2], i = c / 2, d = t.slice([s, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), p = o.slice([s, 0, 0], [a, i, 1]).reshape([1, 1, a, i]), u = n.shape[0], l = n.shape[1], g = T(0, c, 2, "int32"), x = T(1, c, 2, "int32"), $ = ((b) => {
-    const v = b.slice([0, 0, 0, 0], [u, l, a, c]), k = c < r ? b.slice([0, 0, 0, c], [u, l, a, r - c]) : null, h = E(v, g, 3), m = E(v, x, 3), C = h.mul(d), y = m.mul(p), R = C.sub(y), N = m.mul(d), S = h.mul(p), w = N.add(S);
-    h.dispose(), m.dispose(), d.dispose(), p.dispose(), C.dispose(), y.dispose(), N.dispose(), S.dispose();
-    const L = B([R, w], -1);
-    R.dispose(), w.dispose();
-    const P = L.reshape([u, l, a, c]);
-    return L.dispose(), k ? j([P, k], 3) : P;
-  })(n);
-  return g.dispose(), x.dispose(), $;
-}
-function X(o) {
-  const { x: t, sin: e, cos: n } = o.inputs, { pastLen: s } = o.attrs, r = t.shape[3];
-  return V(e, n, r, t, s);
-}
-const Y = {
-  kernelName: "Rope",
-  backendName: "cpu",
-  kernelFunc: X
-};
-f(Y);
-const Z = {
-  kernelName: "Rope",
-  backendName: "tensorflow",
-  kernelFunc: X
-};
-f(Z);
-function st(o, t, e) {
-  return t.ensureRopeCache(o.shape[1] + e), D().runKernel("Rope", { x: o, sin: t.getSin(), cos: t.getCos() }, { pastLen: e });
-}
-const q = {
-  kernelName: "Rope",
-  inputsToSave: ["x", "sin", "cos"],
-  outputsToSave: [],
-  gradFunc: (o, t) => {
-    const [e, n, s] = t, r = n.neg(), c = e.shape[3], i = V(r, s, c, o, 0);
-    return r.dispose(), { x: () => i };
-  }
-};
-A(q);
 export {
-  st as rope,
-  X as ropeCPU
+  s as rope
 };

package/dist/ops/scatterSub.d.ts

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

package/dist/ops/scatterSub.js

@@ -1,150 +1,9 @@
-import { engine as $ } from "@tensorflow/tfjs";
-import { t as u, u as S, v as l, E as f, w as E, o as N, d as h, x as y, r as p, b as x, a as D } from "../index-YPKosni4.js";
-import { c as d } from "../complex-CJ-qCcLB.js";
-import { r as v, s as T } 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 i(e, t = "float32") {
-  if (u(e), t === "complex64") {
-    const a = i(e, "float32"), o = i(e, "float32");
-    return d(a, o);
-  }
-  const r = S(l(e), t);
-  return f.makeTensor(r, e, t);
-}
-/**
- * @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 = "float32") {
-  if (u(e), t === "complex64") {
-    const a = m(e, "float32"), o = i(e, "float32");
-    return d(a, o);
-  }
-  const r = E(l(e), t);
-  return f.makeTensor(r, e, t);
-}
-function C(e, t, r) {
-  const a = t.rank > 1 ? t.shape[t.rank - 1] : 1, o = t.rank > 1 ? t.rank - 1 : 1, s = `Must have updates.shape = indices.shape[:batchDim] + shape[sliceDim:], got updates.shape: ${r.shape}, indices.shape: ${t.shape}, shape: ${e}, sliceDim: ${a}, and batchDim: ${o}.`;
-  if (r.rank < o)
-    throw new Error(s + ` update.rank < ${o}. `);
-  if (e.length < a + (r.rank - o))
-    throw new Error(s + ` Output shape length < ${a + (r.rank - o)}`);
-  if (r.rank !== o + e.length - a)
-    throw new Error(s + ` update.rank != ${o + e.length - a}`);
-  for (let n = 0; n < o; ++n)
-    if (r.shape[n] !== t.shape[n])
-      throw new Error(s + ` updates.shape[${n}] (${r.shape[n]}) != indices.shape[${n}] (${t.shape[n]}).`);
-  for (let n = 0; n < r.rank - o; ++n)
-    if (r.shape[n + o] !== e[n + a])
-      throw new Error(s + ` updates.shape[${n + o}] (${r.shape[n + o]}) != shape[${n + o}] (${e[n + o]})`);
-}
-function O(e, t, r) {
-  if (t.rank < 1)
-    throw new Error(`tf.scatterND() expects the indices to be rank 1 or higher, but the rank was ${t.rank}.`);
-  if (e.rank < 1)
-    throw new Error(`tf.scatterND() expects the updates to be rank 1 or higher, but the rank was ${e.rank}.`);
-  if (t.dtype !== "int32")
-    throw new Error(`The dtype of 'indices' should be int32, but got dtype: ${t.dtype}`);
-  if (r.length < 1)
-    throw new Error(`Output rank must be greater or equal to 1, but got shape: ${r}`);
-  if (r.length === 0) {
-    if (t.size === 0)
-      throw new Error(`Indices specified for empty output. indices shape: ${t.shape}`);
-    if (e.size === 0)
-      throw new Error(`Updates specified for empty output. updates shape: ${e.shape}`);
-  }
-  C(r, t, e);
-}
-/**
- * @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 z(e, t, r) {
-  u(r);
-  const a = h(e, "indices", "scatterND", "int32"), o = h(t, "updates", "scatterND");
-  O(o, a, r);
-  const s = { indices: a, updates: o }, n = { shape: r };
-  return f.runKernel(y, s, n);
-}
-const I = /* @__PURE__ */ N({ scatterND_: z });
-class L {
-  variableNames = ["labels", "softmaxProbs", "dy"];
-  outputShape;
-  userCode;
-  constructor(t, r) {
-    this.outputShape = [t, r], this.userCode = `
-      void main() {
-        ivec2 coords = getOutputCoords();
-        int index = int(getLabels(coords.x));
-        float prob = getSoftmaxProbsAtOutCoords();
-        float dy = getDy(coords.x);
-        setOutput((index == coords.y ? prob - 1.0 : prob) * dy);
-      }
-    `;
-  }
-}
-function P(e) {
-  const { logits: t, labels: r, dy: a } = e.inputs, o = e.backend, s = r.shape[0], n = t.shape[1], c = new L(s, n);
-  return o.runWebGLProgram(c, [r, t, a], "float32");
-}
-const K = {
-  kernelName: "EfficientScatterSub",
-  backendName: "webgl",
-  kernelFunc: P
-};
-p(K);
-function A(e) {
-  const { logits: t, labels: r, dy: a } = e.inputs, o = r.shape[0], s = t.shape[1], n = v(0, o, 1, "int32"), c = T([n, r], 1), b = m([o]), g = I(c, b, [o, s]), k = x(t, g), w = a.reshape([o, 1]);
-  return D(k, w);
-}
-const F = {
-  kernelName: "EfficientScatterSub",
-  backendName: "cpu",
-  kernelFunc: A
-};
-p(F);
-function R(e, t, r) {
-  return $().runKernel("EfficientScatterSub", { logits: e, labels: t, dy: r }, {});
+import { e as i } from "../index-pWA4_lUh.js";
+import "./cpu/scatterSub.js";
+import "./webgl/scatterSub.js";
+function c(t, r, e) {
+  return i().runKernel("EfficientScatterSub", { logits: t, labels: r, dy: e }, {});
 }
 export {
-  R as scatterSub
+  c as scatterSub
 };

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

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

package/dist/ops/webgl/appendCache.js

@@ -0,0 +1,43 @@
+import { r as h } from "../../index-pWA4_lUh.js";
+class m {
+  variableNames = ["cache", "item"];
+  outputShape;
+  userCode;
+  customUniforms = [{ name: "cacheT", type: "int" }];
+  constructor(t, a, o, s, n) {
+    const c = Math.min(o + 1, n);
+    this.outputShape = [t, a, c, s], this.userCode = `
+        void main() {
+            ivec4 coords = getOutputCoords(); // [b, h, t, d]
+            int b = coords.x;
+            int h = coords.y;
+            int t = coords.z;
+            int d = coords.w;
+
+            int itemT = 1;
+            int maxSize = ${n};
+            int totalT = cacheT + itemT;
+            int start = totalT >= maxSize ? 1 : 0;
+
+            int srcT = t + start;
+            float val = 0.0;
+            if (srcT < cacheT) {
+                val = getCache(b, h, srcT, d);
+            } else {
+                val = getItem(b, h, 0, d);
+            }
+            setOutput(val);
+        }
+        `;
+  }
+}
+function p(e) {
+  const { cache: t, item: a } = e.inputs, { maxSize: o } = e.attrs, s = e.backend, n = t.shape[0], c = t.shape[2], r = t.shape[1], i = new m(n, r, c, a.shape[3], o);
+  return s.runWebGLProgram(i, [t, a], "float32", [[c]]);
+}
+const d = {
+  kernelName: "AppendCache",
+  backendName: "webgl",
+  kernelFunc: p
+};
+h(d);

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

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

package/dist/ops/webgl/attentionMask.js

@@ -0,0 +1,43 @@
+import { r as m } from "../../index-pWA4_lUh.js";
+class k {
+  variableNames = ["q", "k", "mask"];
+  outputShape;
+  userCode;
+  customUniforms = [{ name: "divisor", type: "float" }];
+  constructor(t, a, o, e) {
+    this.outputShape = [t, a, o, o], 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 < ${e}; ++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);
+        }
+        `;
+  }
+}
+function l(s) {
+  const { q: t, k: a, mask: o } = s.inputs, { divisor: e } = s.attrs, n = s.backend, i = t.shape[0], r = t.shape[2], c = t.shape[1], d = new k(i, c, r, t.shape[3]);
+  return n.runWebGLProgram(d, [t, a, o], "float32", [[e]]);
+}
+const u = {
+  kernelName: "AttentionMask",
+  backendName: "webgl",
+  kernelFunc: l
+};
+m(u);

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

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

package/dist/ops/webgl/gatherSub.js

@@ -0,0 +1,27 @@
+import { r as l } from "../../index-pWA4_lUh.js";
+class u {
+  variableNames = ["labels", "logits", "values"];
+  outputShape;
+  userCode;
+  constructor(e) {
+    this.outputShape = [e], this.userCode = `
+      void main() {
+        int coords = getOutputCoords();
+        int index = int(getLabelsAtOutCoords());
+        float val = getValuesAtOutCoords();
+        float logit = getLogits(coords, index);
+        setOutput(val - logit);
+      }
+    `;
+  }
+}
+function i(t) {
+  const { logits: e, labels: o, values: s } = t.inputs, a = t.backend, r = o.shape[0], n = new u(r);
+  return a.runWebGLProgram(n, [o, e, s], "float32");
+}
+const c = {
+  kernelName: "EfficientGatherSub",
+  backendName: "webgl",
+  kernelFunc: i
+};
+l(c);

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

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

package/dist/ops/webgl/qkv.js

@@ -0,0 +1,46 @@
+import { r as i } from "../../index-pWA4_lUh.js";
+class l {
+  variableNames = ["x", "kernel"];
+  outputShape;
+  userCode;
+  // enableShapeUniforms = true;
+  customUniforms = [{ name: "mode", type: "int" }];
+  constructor(e, t, s, o) {
+    const n = o / t;
+    this.outputShape = [e, t, s, n], this.userCode = `
+        void main() {
+            ivec4 coords = getOutputCoords(); // [b, h, t, d]
+            int b = coords.x;
+            int h = coords.y;
+            int t = coords.z;
+            int d = coords.w;
+
+            // Compute output channel index in fused kernel
+            int out_offset = mode * ${t} * ${n} + h * ${n} + d;
+
+            float sum = 0.0;
+            for (int c = 0; c < ${o}; ++c) {
+                float xval = getX(b, t, c); // fetch from x
+                float kval = getKernel(c, out_offset); // fetch from kernel
+                sum += xval * kval;
+            }
+
+            setOutput(sum);
+        }
+        `;
+  }
+}
+function m(r) {
+  const { x: e, kernel: t } = r.inputs, { heads: s } = r.attrs, o = r.backend, n = e.shape[0], c = e.shape[1], u = e.shape[2], a = new l(n, s, c, u);
+  return [
+    o.runWebGLProgram(a, [e, t], "float32", [[0]]),
+    o.runWebGLProgram(a, [e, t], "float32", [[1]]),
+    o.runWebGLProgram(a, [e, t], "float32", [[2]])
+  ];
+}
+const d = {
+  kernelName: "QKV",
+  backendName: "webgl",
+  kernelFunc: m
+};
+i(d);

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

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