@genai-fi/nanogpt 0.9.0 → 0.10.0

package/dist/ops/webgpu/attentionMask.js

@@ -1,71 +1,28 @@
-import { f, a4 as m } from "../../index-BzFyqcy-.js";
-import { g as k } from "../../webgpu_program-DkQJOJSd.js";
-import { f as l, c as v } from "../../webgpu_util-pLEV9tks.js";
-class g {
-  variableNames = ["q", "k"];
-  outputShape;
-  shaderKey = "AttentionMask";
-  dispatchLayout;
-  dispatch;
-  uniforms = "divisor: f32, pastLen: i32, inf: f32";
-  workgroupSize = [64, 1, 1];
-  size = !0;
-  hs;
-  nh;
-  T1;
-  T2;
-  constructor(t, e, s, o, i) {
-    if (this.shaderKey = `AttentionMask_${i}`, this.outputShape = [t, e, s, o], this.hs = i, this.nh = e, this.T1 = s, this.T2 = o, this.dispatchLayout = l(this.outputShape), this.dispatch = v(this.dispatchLayout, this.outputShape, this.workgroupSize), i % 4 !== 0)
-      throw new Error("Head size must be a multiple of 4 for AttentionMaskProgram");
-  }
-  getUserCode() {
-    return `
-            ${k("index")} {
-                
-                let coords = getCoordsFromIndex(index);
-                let b = coords[0];
-                let h = coords[1];
-                let t1 = coords[2];
-                let t2 = coords[3];
-
-                if (index < uniforms.size) {
-                    if (t2 > t1 + uniforms.pastLen) {
-                        setOutputAtIndex(index, uniforms.inf);
-                        return;
-                    }
-
-                    let q0 = getIndexFromCoords4D(vec4<i32>(b, h, t1, 0), uniforms.qShape);
-                    let k0 = getIndexFromCoords4D(vec4<i32>(b, h, t2, 0), uniforms.kShape);
-                    
-                    var sum: f32 = 0.0;
-                    for (var i: i32 = 0; i < ${this.hs}; i = i + 4) {
-                        let qv = vec4<f32>(q[q0 + i], q[q0 + i + 1], q[q0 + i + 2], q[q0 + i + 3]);
-                        let kv = vec4<f32>(k[k0 + i], k[k0 + i + 1], k[k0 + i + 2], k[k0 + i + 3]);
-                        sum = sum + dot(qv, kv);
-                    }
-                    let scaled = sum * uniforms.divisor;
-                    setOutputAtIndex(index, scaled);
-                }
-            }
-        `;
-  }
-}
-function b(n) {
-  const { q: t, k: e } = n.inputs, { divisor: s, pastLen: o } = n.attrs, i = n.backend, r = t.shape[0], p = t.shape[2], a = e.shape[2], u = t.shape[1], h = t.shape[3];
-  if (m(e.shape, [r, u, a, h], "Error in AttentionMask: "), s === 0)
+import "../../index-ZyQhjEPo.js";
+import { j as d } from "../../tensor-DdQUJZlz.js";
+import { isPackedTensor as p } from "../../utilities/packed.js";
+import { b } from "../../matMul16--R5hOwDG.js";
+import l from "./attentionMask32_program.js";
+import { r as M } from "../../tensor_util-DV-FP5Q3.js";
+function w(n) {
+  const { q: t, k: e } = n.inputs, { divisor: a, pastLen: o } = n.attrs, m = n.backend;
+  if (p(t) && p(e))
+    return b(t, e, !1, !0, { causalMask: !0, pastLen: o, scale: a });
+  const r = t.shape[0], k = t.shape[2], s = e.shape[2], i = t.shape[1], c = t.shape[3];
+  if (d(e.shape, [r, i, s, c], "Error in AttentionMask: "), a === 0)
     throw new Error("Divisor must be non-zero in AttentionMask");
   if (o < 0)
     throw new Error("pastLen must be non-negative in AttentionMask");
-  const c = new g(r, u, p, a, h), d = [
-    { type: "float32", data: [s] },
+  const u = new l(r, i, k, s, c), f = [
+    { type: "float32", data: [a] },
     { type: "int32", data: [o] },
     { type: "float32", data: [Number.NEGATIVE_INFINITY] }
-  ];
-  return i.runWebGPUProgram(c, [t, e], "float32", d);
+  ], h = t.dtype;
+  return m.runWebGPUProgram(u, [t, e], h, f);
 }
-const q = {
+const A = {
   kernelName: "AttentionMask",
   backendName: "webgpu",
-  kernelFunc: b
+  kernelFunc: w
 };
-f(q);
+M(A);

package/dist/ops/webgpu/attentionMask32_program.d.ts

@@ -0,0 +1,19 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class AttentionMaskProgram32 implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    uniforms: string;
+    workgroupSize: [number, number, number];
+    size: boolean;
+    hs: number;
+    nh: number;
+    T1: number;
+    T2: number;
+    constructor(batch: number, nh: number, T1: number, T2: number, hs: number);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/attentionMask32_program.js

@@ -0,0 +1,54 @@
+import { e as r } from "../../webgpu_program-Cigz-7RF.js";
+import { f as a, c as u } from "../../webgpu_util-BBCnKm2X.js";
+class p {
+  variableNames = ["q", "k"];
+  outputShape;
+  shaderKey = "AttentionMask";
+  dispatchLayout;
+  dispatch;
+  uniforms = "divisor: f32, pastLen: i32, inf: f32";
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  hs;
+  nh;
+  T1;
+  T2;
+  constructor(e, i, s, o, t) {
+    if (this.shaderKey = `AttentionMask_${t}`, this.outputShape = [e, i, s, o], this.hs = t, this.nh = i, this.T1 = s, this.T2 = o, this.dispatchLayout = a(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize), t % 4 !== 0)
+      throw new Error("Head size must be a multiple of 4 for AttentionMaskProgram");
+  }
+  getUserCode() {
+    return `
+            ${r("index")} {
+                
+                let coords = getCoordsFromIndex(index);
+                let b = coords[0];
+                let h = coords[1];
+                let t1 = coords[2];
+                let t2 = coords[3];
+
+                if (index < uniforms.size) {
+                    if (t2 > t1 + uniforms.pastLen) {
+                        setOutputAtIndex(index, uniforms.inf);
+                        return;
+                    }
+
+                    let q0 = getIndexFromCoords4D(vec4<i32>(b, h, t1, 0), uniforms.qShape);
+                    let k0 = getIndexFromCoords4D(vec4<i32>(b, h, t2, 0), uniforms.kShape);
+                    
+                    var sum: f32 = 0.0;
+                    for (var i: i32 = 0; i < ${this.hs}; i = i + 4) {
+                        let qv = vec4<f32>(q[q0 + i], q[q0 + i + 1], q[q0 + i + 2], q[q0 + i + 3]);
+                        let kv = vec4<f32>(k[k0 + i], k[k0 + i + 1], k[k0 + i + 2], k[k0 + i + 3]);
+                        sum = sum + dot(qv, kv);
+                    }
+                    let scaled = sum * uniforms.divisor;
+                    setOutputAtIndex(index, scaled);
+                }
+            }
+        `;
+  }
+}
+export {
+  p as default
+};

package/dist/ops/webgpu/concat16.d.ts

@@ -0,0 +1,19 @@
+import { KernelConfig } from '@tensorflow/tfjs-core';
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu_program';
+export declare class ConcatProgram implements WebGPUProgram {
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    variableNames: string[];
+    uniforms: string;
+    workPerThread: number;
+    workgroupSize: [number, number, number];
+    size: boolean;
+    offsetLength: number;
+    constructor(shapes: Array<[number, number]>);
+    getUserCode(): string;
+}
+export declare const concatConfig: KernelConfig;

package/dist/ops/webgpu/concat16.js

@@ -0,0 +1,128 @@
+import "../../index-ZyQhjEPo.js";
+import { e as x } from "../../webgpu_program-Cigz-7RF.js";
+import { f as I, c as D } from "../../webgpu_util-BBCnKm2X.js";
+import { r as y } from "../../Reshape-CF6odzV4.js";
+import { r as $ } from "../../tensor_util-DV-FP5Q3.js";
+import { p as F, s as c } from "../../tensor-DdQUJZlz.js";
+import { a as L, c as d } from "../../concat_util-DcJk7YHS.js";
+class T {
+  outputShape;
+  shaderKey;
+  dispatchLayout;
+  dispatch;
+  variableNames;
+  uniforms = "";
+  workPerThread = 1;
+  workgroupSize = [64, 1, 1];
+  size = !0;
+  offsetLength;
+  constructor(t) {
+    this.outputShape = d(
+      t,
+      1
+      /* axis */
+    ), this.variableNames = t.map((e, a) => `T${a}`), this.dispatchLayout = I(this.outputShape), this.dispatch = D(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+      this.workPerThread,
+      1,
+      1
+    ]), this.offsetLength = t.length - 1;
+    for (let e = 0; e < this.offsetLength; e++)
+      this.uniforms += `offset${e} : i32,`;
+    this.shaderKey = "concat16";
+  }
+  getUserCode() {
+    const t = [];
+    if (this.offsetLength > 0) {
+      t.push(
+        "if (yC < uniforms.offset0){ result[getIndexFromCoords2D(coords, uniforms.outShape)] = T0[getIndexFromCoords2D(vec2<i32>(yR, yC), uniforms.t0Shape)]; }"
+      );
+      for (let s = 1; s < this.offsetLength; s++)
+        t.push(
+          `else if (yC < uniforms.offset${[s]}){ result[getIndexFromCoords2D(coords, uniforms.outShape)] = T${s}[getIndexFromCoords2D(vec2<i32>(yR, yC - uniforms.offset${s - 1}), uniforms.t${s}Shape)]; }`
+        );
+      const a = this.offsetLength, i = this.offsetLength - 1;
+      t.push(
+        `else { result[getIndexFromCoords2D(coords, uniforms.outShape)] = T${a}[getIndexFromCoords2D(vec2<i32>(yR, yC - uniforms.offset${i}), uniforms.t${a}Shape)]; }`
+      );
+    } else
+      t.push(
+        "result[getIndexFromCoords2D(coords, uniforms.outShape)] = T0[getIndexFromCoords2D(vec2<i32>(yR, yC), uniforms.t0Shape)];"
+      );
+    return `
+      ${x("index")} {
+        for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
+          let flatIndex = index * ${this.workPerThread} + i;
+          if(flatIndex < uniforms.size) {
+            let coords = getCoordsFromIndex(flatIndex);
+            let yR = coords.x;
+            let yC = coords.y;
+
+            ${t.join(`
+        `)}
+          }
+        }
+      }
+    `;
+  }
+}
+function m(n, t, e) {
+  const a = e.device.limits.maxStorageBuffersPerShaderStage - 1;
+  if (n.length > a) {
+    const o = [];
+    for (let p = 0; p < n.length; p += a) {
+      const C = n.slice(p, p + a);
+      o.push(m(C, t, e));
+    }
+    const S = m(o, t, e);
+    for (const p of o)
+      e.disposeData(p.dataId);
+    return S;
+  }
+  const { tensors2D: i, outShape: s } = P(n, t, e), h = i.map((o) => o.shape), u = new T(h), f = [], r = new Array(h.length - 1);
+  if (r.length > 0) {
+    r[0] = h[0][1], f.push({ type: "int32", data: [r[0]] });
+    for (let o = 1; o < r.length; o++)
+      r[o] = r[o - 1] + h[o][1], f.push({ type: "int32", data: [r[o]] });
+  }
+  const l = e.runWebGPUProgram(u, i, i[0].dtype, f);
+  i.forEach((o) => e.disposeData(o.dataId));
+  const g = y({ inputs: { x: l }, backend: e, attrs: { shape: s } });
+  return e.disposeData(l.dataId), g.packed = !0, g;
+}
+function P(n, t, e) {
+  const a = d(
+    n.map((s) => s.shape),
+    t
+  );
+  return { tensors2D: n.map(
+    (s) => y({
+      inputs: { x: s },
+      backend: e,
+      attrs: {
+        shape: [c(s.shape.slice(0, t)), c(s.shape.slice(t))]
+      }
+    })
+  ), outShape: a };
+}
+function w(n) {
+  const { inputs: t, backend: e, attrs: a } = n, { axis: i } = a, s = F(i, t[0].shape)[0], h = t.map((r) => r.shape);
+  L(h, s);
+  const u = d(
+    t.map((r) => r.shape),
+    s
+  );
+  if (c(u) === 0)
+    return e.makeTensorInfo(u, t[0].dtype, []);
+  const f = t.filter((r) => c(r.shape) > 0);
+  return m(f, s, e);
+}
+const v = {
+  kernelName: "Concat16",
+  backendName: "webgpu",
+  kernelFunc: w
+};
+$(v);
+export {
+  T as ConcatProgram,
+  v as concatConfig
+};

package/dist/ops/webgpu/gatherSub.js

@@ -1,6 +1,8 @@
-import { g as u } from "../../webgpu_program-DkQJOJSd.js";
-import { f as h, c as p } from "../../webgpu_util-pLEV9tks.js";
-import { f as c, a4 as r } from "../../index-BzFyqcy-.js";
+import { e as u } from "../../webgpu_program-Cigz-7RF.js";
+import { f as p, c as h } from "../../webgpu_util-BBCnKm2X.js";
+import "../../index-ZyQhjEPo.js";
+import { j as s } from "../../tensor-DdQUJZlz.js";
+import { r as c } from "../../tensor_util-DV-FP5Q3.js";
 class l {
   variableNames = ["labels", "logits", "values"];
   outputShape;
@@ -10,7 +12,7 @@ class l {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(e) {
-    this.outputShape = [e], this.dispatchLayout = h(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = [e], this.dispatchLayout = p(this.outputShape), this.dispatch = h(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
@@ -26,10 +28,10 @@ class l {
   }
 }
 function d(t) {
-  const { logits: e, labels: a, values: s } = t.inputs, o = t.backend, i = a.shape[0];
-  r(s.shape, [i], "Error in EfficientGatherSub: "), r(a.shape, [i], "Error in EfficientGatherSub: ");
+  const { logits: e, labels: a, values: r } = t.inputs, o = t.backend, i = a.shape[0];
+  s(r.shape, [i], "Error in EfficientGatherSub: "), s(a.shape, [i], "Error in EfficientGatherSub: ");
   const n = new l(i);
-  return o.runWebGPUProgram(n, [a, e, s], "float32");
+  return o.runWebGPUProgram(n, [a, e, r], "float32");
 }
 const f = {
   kernelName: "EfficientGatherSub",

package/dist/ops/webgpu/gelu.js

@@ -1,8 +1,10 @@
-import { f as s } from "../../index-BzFyqcy-.js";
-import { g as a } from "../../webgpu_program-DkQJOJSd.js";
-import { f as o, c as p } from "../../webgpu_util-pLEV9tks.js";
-const u = 0.7978845608028654, i = 0.044715;
-class d {
+import "../../index-ZyQhjEPo.js";
+import { e as s } from "../../webgpu_program-Cigz-7RF.js";
+import { f as o, c as p } from "../../webgpu_util-BBCnKm2X.js";
+import { isPackedTensor as l } from "../../utilities/packed.js";
+import { r as h } from "../../tensor_util-DV-FP5Q3.js";
+const r = 0.7978845608028654, u = 0.044715;
+class x {
   outputShape;
   shaderKey;
   dispatchLayout;
@@ -21,13 +23,13 @@ class d {
             }
             fn unaryOperation(x : f32) -> f32 {
                 let x3 = x * x * x;
-                var inner = fma(${i}, x3, x);
-                inner = ${u} * inner;
+                var inner = fma(${u}, x3, x);
+                inner = ${r} * inner;
                 inner = tanhComplete(inner);
                 inner = 0.5 * (1.0 + inner);
                 return x * inner;
             }
-            ${a("index")} {
+            ${s("index")} {
                 if (index < uniforms.size) {
                 let a = getAByOutputIndex(index);
                 setOutputAtIndex(index, unaryOperation(a));
@@ -36,17 +38,58 @@ class d {
       `;
   }
 }
-function c(t) {
-  const { x: e } = t.inputs, n = t.backend, r = new d(e.shape);
-  return n.runWebGPUProgram(r, [e], "float32");
+function g(t) {
+  const { x: e } = t.inputs, a = t.backend, i = new x(e.shape);
+  return a.runWebGPUProgram(i, [e], "float32");
 }
-const l = {
+const f = {
   kernelName: "Gelu",
   backendName: "webgpu",
-  kernelFunc: c
+  kernelFunc: g
 };
-s(l);
-class x {
+h(f);
+class m {
+  // Inputs: dy, x
+  variableNames = ["dy", "x"];
+  outputShape;
+  shaderKey = "GeluGrad";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [128, 1, 1];
+  size = !0;
+  constructor(e) {
+    this.outputShape = e, this.dispatchLayout = o(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  }
+  getUserCode() {
+    return `
+            // TODO: revisit after https://github.com/gpuweb/gpuweb/issues/4458 is resolved
+            fn tanhComplete(x: f32) -> f32 {
+                return select(tanh(x), sign(x), abs(x) > 15.0);
+            }
+            fn activationGrad(dy: f32, X: f32) -> f32 {
+                let x2 = X * X;
+                let x3 = x2 * X;
+                let u  = ${r} * (X + ${u} * x3);
+                let t  = tanhComplete(u);
+                let sech2 = 1.0 - t * t;
+                let du_dx = ${r} * (1.0 + 3.0 * ${u} * x2);
+                let dgelu = 0.5 * (1.0 + t) + 0.5 * X * sech2 * du_dx;
+                return dy *dgelu;
+            }
+            ${s("index")} {
+                if (index < uniforms.size) {
+                    let X  = unpack2x16float(u32(x[index]));
+                    let DY = unpack2x16float(u32(dy[index]));
+                    let dgelu = vec2<f32>(
+                        activationGrad(DY.x, X.x),
+                        activationGrad(DY.y, X.y)
+                    );
+                    result[index] = i32(pack2x16float(dgelu));
+                }
+            }`;
+  }
+}
+class y {
   // Inputs: dy, x
   variableNames = ["dy", "x"];
   outputShape;
@@ -64,32 +107,36 @@ class x {
             fn tanhComplete(x: f32) -> f32 {
                 return select(tanh(x), sign(x), abs(x) > 15.0);
             }
-            ${a("index")} {
+            fn activationGrad(dy: f32, X: f32) -> f32 {
+                let x2 = X * X;
+                let x3 = x2 * X;
+                let u  = ${r} * (X + ${u} * x3);
+                let t  = tanhComplete(u);
+                let sech2 = 1.0 - t * t;
+                let du_dx = ${r} * (1.0 + 3.0 * ${u} * x2);
+                let dgelu = 0.5 * (1.0 + t) + 0.5 * X * sech2 * du_dx;
+                return dy *dgelu;
+            }
+            ${s("index")} {
                 if (index < uniforms.size) {
                     let X  = getXByOutputIndex(index);
-                    let x2 = X * X;
-                    let x3 = x2 * X;
-                    let u  = ${u} * (X + ${i} * x3);
-                    let t  = tanhComplete(u);
-                    let sech2 = 1.0 - t * t;
-                    let du_dx = ${u} * (1.0 + 3.0 * ${i} * x2);
-                    let dgelu = 0.5 * (1.0 + t) + 0.5 * X * sech2 * du_dx;
                     let DY = getDyByOutputIndex(index);
-                    setOutputAtIndex(index, DY * dgelu);
+                    let dgelu = activationGrad(DY, X);
+                    setOutputAtIndex(index, dgelu);
                 }
             }`;
   }
 }
-function g(t) {
-  const { dy: e, x: n } = t.inputs, r = t.backend, h = new x(n.shape);
-  return r.runWebGPUProgram(h, [e, n], "float32");
+function b(t) {
+  const { dy: e, x: a } = t.inputs, i = t.backend, n = l(e), c = n ? new m(a.shape) : new y(a.shape), d = i.runWebGPUProgram(c, [e, a], n ? "int32" : "float32");
+  return d.packed = n, d;
 }
-const f = {
+const k = {
   kernelName: "GeluGrad",
   backendName: "webgpu",
-  kernelFunc: g
+  kernelFunc: b
 };
-s(f);
+h(k);
 export {
-  d as GeluProgram
+  x as GeluProgram
 };

package/dist/ops/webgpu/index.js

@@ -9,3 +9,15 @@ import "./qkv.js";
 import "./gelu.js";
 import "./adamMoments.js";
 import "./adamAdjust.js";
+import "./pack16.js";
+import "./unpack16.js";
+import "./softmax16.js";
+import "./matMul16.js";
+import "./transpose16.js";
+import "./sum16.js";
+import "./slice16.js";
+import "./add16.js";
+import "./concat16.js";
+import "./mul16.js";
+import "./sub16.js";
+import "./softmax16grad.js";

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

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

package/dist/ops/webgpu/matMul16.js

@@ -0,0 +1,58 @@
+import { m as y, b as B, j as Q } from "../../index-ZyQhjEPo.js";
+import { isPackedTensor as R } from "../../utilities/packed.js";
+import { reshape16 as U } from "../reshape16.js";
+import { matMulMul as V } from "../matMulMul.js";
+import { matMulGelu as X } from "../matMulGelu.js";
+import Y from "./matMul16_program.js";
+import { r as Z } from "../../tensor_util-DV-FP5Q3.js";
+import { m as _ } from "../../mat_mul-DeAh4uTH.js";
+import { r as x } from "../../reshape-DevtBWtf.js";
+import { t as C } from "../../transpose-DKELTqhe.js";
+import { s as E } from "../../tensor-DdQUJZlz.js";
+function $(p) {
+  const { A: e, B: s } = p.inputs, { transposeA: d, transposeB: f, scale: i, activation: k, scaleA: c, scaleB: u, forceOutputShape: o, perm: m, causalMask: g, pastLen: W } = p.attrs, z = p.backend, S = !R(e), M = !R(s);
+  if (S && M) {
+    const A = c !== void 0 ? y(e, B(c)) : e, b = u !== void 0 ? y(s, B(u)) : s;
+    if (g)
+      throw new Error("Causal mask is not supported for unpacked MatMul16.");
+    let a;
+    if (i !== void 0 ? a = V(A, b, B(i), d, f) : k === "gelu" ? a = X(A, b) : a = _(A, b, d, f), m)
+      if (o) {
+        const n = x(a, o);
+        a.dispose();
+        const J = C(n, m);
+        return n.dispose(), J;
+      } else {
+        const n = C(a, m);
+        return a.dispose(), n;
+      }
+    else if (o) {
+      const n = x(a, o);
+      return a.dispose(), n;
+    } else
+      return a;
+  }
+  if (S && !M)
+    throw new Error("When using mixed precision, A must be packed if B is packed.");
+  if (!S && M)
+    throw new Error("When using mixed precision, B must be packed if A is packed.");
+  const h = e.shape.length, l = s.shape.length, F = e.shape.slice(0, -2), I = s.shape.slice(0, -2), v = E(F), w = E(I), N = Q(e.shape.slice(0, -2), s.shape.slice(0, -2)), j = Math.max(v, w), K = e.shape[h - 2], L = s.shape[l - 2], T = e.shape[h - 1] * 2, q = s.shape[l - 1] * 2, D = U(e, [v, e.shape[h - 2], e.shape[h - 1]]), G = U(s, [w, s.shape[l - 2], s.shape[l - 1]]), t = new Y(j, K, L, T, q, d, f), r = [];
+  i !== void 0 && (t.useScale(), r.push({ type: "float32", data: [i] })), c !== void 0 && (t.useScaleA(), r.push({ type: "float32", data: [c] })), u !== void 0 && (t.useScaleB(), r.push({ type: "float32", data: [u] })), k !== void 0 && t.useActivation(k), g && (t.useCausalMask(), r.push({ type: "int32", data: [W || 0] }));
+  const O = t.outputShape.length;
+  o && (p.attrs.originalShape = t.outputShape);
+  const H = o ?? N.concat([t.outputShape[O - 2], t.outputShape[O - 1]]);
+  t.setOutputShape(H, m);
+  const P = z.runWebGPUProgram(
+    t,
+    [D, G],
+    "int32",
+    r.length > 0 ? r : void 0
+  );
+  return P.packed = !0, D.dispose(), G.dispose(), P;
+}
+const ee = {
+  kernelName: "MatMul16",
+  backendName: "webgpu",
+  kernelFunc: $
+};
+Z(ee);

package/dist/ops/webgpu/matMul16_program.d.ts

@@ -0,0 +1,42 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class MatMul16ProgramGeneric implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+        y: number[];
+        z: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    dimInner: number;
+    transposeA: boolean;
+    transposeB: boolean;
+    broadcastBatch: boolean;
+    tileInner: number;
+    uniforms?: string;
+    scale: boolean;
+    scaleA: boolean;
+    scaleB: boolean;
+    activation?: 'gelu';
+    causalMask: boolean;
+    outputComponent?: number | undefined;
+    variableComponents?: number[];
+    outputIndexSnippet?: string;
+    outputStrideSnippet?: string;
+    constructor(batch: number, O1: number, O2: number, I1: number, I2: number, transposeA?: boolean, transposeB?: boolean);
+    private addUniform;
+    setOutputShape(shape: number[], perm?: number[]): void;
+    useScale(): void;
+    useScaleA(): void;
+    useScaleB(): void;
+    useActivation(activation: 'gelu'): void;
+    useCausalMask(): void;
+    private activationSnippet;
+    private readASnippet;
+    private readBSnippet;
+    private baseIndexSnippets;
+    private offsetSnippets;
+    getUserCode(): string;
+}