@genai-fi/nanogpt 0.6.3 → 0.7.1

package/dist/webgpu_program-WaoMq-WD.js

@@ -0,0 +1,548 @@
+import { aa as k, ab as z, ac as E, a1 as j, l as A } from "./index-UdZhlibC.js";
+/**
+ * @license
+ * Copyright 2019 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 L(e, s) {
+  if (Math.max(...e) > 5)
+    throw new Error("Cannot symbolically compute strides for rank > 6 tensor.");
+  const t = e.length, o = "xyzwuv", i = e.map((n) => `${s}.${o[n]}`), u = new Array(t - 1);
+  u[t - 2] = i[t - 1];
+  for (let n = t - 3; n >= 0; --n)
+    u[n] = `(${u[n + 1]} * ${i[n + 1]})`;
+  return u;
+}
+const X = (e, s, t) => t === "int32" ? `atomicAdd(${e}, bitcast<i32>(${s}));` : `
+          {
+            var oldValue = 0;
+            loop {
+              let newValueF32 = bitcast<f32>(oldValue) + (${s});
+              let newValue = bitcast<i32>(newValueF32);
+              let res = atomicCompareExchangeWeak(${e}, oldValue, newValue);
+              if res.exchanged {
+                break;
+              }
+              oldValue = res.old_value;
+            }
+          }`;
+/**
+ * @license
+ * Copyright 2022 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.
+ * =============================================================================
+ */
+var b;
+(function(e) {
+  e[e.FROM_PIXELS = 0] = "FROM_PIXELS", e[e.DRAW = 1] = "DRAW";
+})(b || (b = {}));
+const H = (e, s, t, o, i) => {
+  const u = { dtype: o.dtype, shape: o.shape }, n = D(t, u, s), r = e.createShaderModule({ code: n, label: s.constructor.name });
+  let d = E().get("WEBGPU_PRINT_SHADER");
+  if (d !== "") {
+    d = d.toLowerCase();
+    const p = d.split(",");
+    (d === "all" || p.some((a) => s.shaderKey.toLowerCase().includes(a))) && (console.group(s.shaderKey), console.debug(n), console.groupEnd());
+  }
+  return i ? e.createComputePipelineAsync({
+    compute: { module: r, entryPoint: "_start" },
+    label: s.constructor.name,
+    layout: "auto"
+  }) : e.createComputePipeline({
+    compute: { module: r, entryPoint: "_start" },
+    label: s.constructor.name,
+    layout: "auto"
+  });
+}, f = (e, s = "f32") => {
+  switch (e) {
+    case 1:
+      return `${s}`;
+    case 2:
+      return `vec2<${s}>`;
+    case 3:
+      return `vec3<${s}>`;
+    case 4:
+      return `vec4<${s}>`;
+    default:
+      throw new Error(`${e}-component ${s} is not supported.`);
+  }
+};
+function v(e) {
+  if (e <= 1)
+    return "i32";
+  if (e === 2)
+    return "vec2<i32>";
+  if (e === 3)
+    return "vec3<i32>";
+  if (e === 4)
+    return "vec4<i32>";
+  if (e === 5)
+    return "vec5";
+  if (e === 6)
+    return "vec6";
+  throw Error(`GPU for rank ${e} is not yet supported`);
+}
+function I(e) {
+  if (e === 0)
+    return "x";
+  if (e === 1)
+    return "y";
+  if (e === 2)
+    return "z";
+  if (e === 3)
+    return "w";
+  if (e === 4)
+    return "u";
+  if (e === 5)
+    return "v";
+  throw Error(`Index ${e} is not yet supported`);
+}
+function Y(...e) {
+  let s;
+  switch (e.length) {
+    case 0:
+      s = `
+        fn main()
+      `;
+      break;
+    case 1:
+      s = `
+        fn main(${e[0]} : i32)
+      `;
+      break;
+    default:
+      throw Error("Unreachable");
+  }
+  return s;
+}
+function w(e, s) {
+  let t;
+  return t = `
+     ${N(s)}
+      fn _start(@builtin(local_invocation_id) LocalId : vec3<u32>,
+                @builtin(global_invocation_id) GlobalId : vec3<u32>,
+                @builtin(local_invocation_index) LocalIndex: u32,
+                @builtin(workgroup_id) WorkgroupId : vec3<u32>,
+                @builtin(num_workgroups) NumWorkgroups : vec3<u32>) {
+        localId = LocalId;
+        localIndex = LocalIndex;
+        globalId = GlobalId;
+        numWorkgroups = NumWorkgroups;
+        workgroupId = WorkgroupId;
+        ${e ? "main(getGlobalIndex());" : "main();"};
+      }
+    `, t;
+}
+function N(e) {
+  return `
+  @compute @workgroup_size(${e.workgroupSize[0]}, ${e.workgroupSize[1]}, ${e.workgroupSize[2]})
+`;
+}
+function D(e, s, t) {
+  const o = [], i = t.workgroupSize[0] * t.workgroupSize[1] * t.workgroupSize[2];
+  if (t.outputComponent = t.outputComponent ? t.outputComponent : 1, o.push(`
+
+      var<private> localId: vec3<u32>;
+      var<private> localIndex: u32;
+      var<private> globalId: vec3<u32>;
+      var<private> numWorkgroups: vec3<u32>;
+      var<private> workgroupId: vec3<u32>;
+
+      // Only used when the y/z dimension of workgroup size is 1.
+      fn getGlobalIndex() -> i32 {
+        ${F(t) ? "  return i32(globalId.x);" : `  return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +
+                workgroupId.y * numWorkgroups.x + workgroupId.x) * ${i}u +
+                localIndex);
+        `}
+      }
+    `), t.pixelsOpType != null) {
+    const h = t.pixelsOpType === b.FROM_PIXELS ? `@group(0) @binding(0) var<storage, read_write> result: array<${x(s.dtype, t.outputComponent)}>;` : `@group(0) @binding(1) var<storage, read> inBuf : array<${x(e[0].dtype, t.outputComponent)}>;`, c = s.shape.length === 3 ? "vec2<i32>" : "i32";
+    o.push(`
+        struct Uniform {
+          outShapeStrides : ${c},
+          size            : i32,
+          numChannels     : i32,
+          alpha           : f32,
+        };
+
+        ${h}
+        @group(0) @binding(2) var<uniform> uniforms: Uniform;
+      `);
+    const S = m(t);
+    return [
+      C,
+      o.join(`
+`),
+      y(s.shape),
+      t.getUserCode(),
+      w(S, t)
+    ].join(`
+`);
+  }
+  let u, n, r = "struct Uniforms { NAN : f32, INFINITY : f32, ";
+  t.variableNames.forEach((h, c) => {
+    const S = v(e[c].shape.length);
+    r += `${h.charAt(0).toLowerCase() + h.slice(1)}Shape : ${S}, `, u = e[c].shape.length - 1, n = v(u), r += `${h.charAt(0).toLowerCase() + h.slice(1)}ShapeStrides: ${n}, `;
+  });
+  const d = v(s.shape.length);
+  r += `outShape : ${d}, `, u = s.shape.length - 1, n = v(u), r += `
+         outShapeStrides: ${n}, `, t.size && (r += "size : i32, "), t.uniforms && (r += t.uniforms), r += "};", r = B(r), o.push(r), t.atomic ? o.push(`
+      @group(0) @binding(0) var<storage, read_write> result: array<atomic<i32>>;
+    `) : o.push(`
+      @group(0) @binding(0) var<storage, read_write> result: array<${x(s.dtype, t.outputComponent)}>;
+    `), t.variableNames.forEach((h, c) => {
+    o.push(`
+      @group(0) @binding(${1 + c}) var<storage, read> ${h}: array<${t.variableComponents ? x(e[c].dtype, t.variableComponents[c]) : x(e[c].dtype, t.outputComponent)}>;
+        `);
+  }), r !== "" && o.push(`
+      @group(0) @binding(${1 + t.variableNames.length}) var<uniform> uniforms: Uniforms;
+      `);
+  const p = R(s.shape, t.dispatchLayout), a = [
+    C,
+    o.join(`
+`) + T,
+    y(s.shape),
+    p,
+    G(s.shape.length)
+  ];
+  t.atomic || a.push(V(s.shape, s.dtype, t.outputComponent)), t.variableNames.forEach((h, c) => {
+    a.push(`${y(e[c].shape, h)}`);
+  });
+  const g = e.map((h, c) => P(h, s.shape, t.variableComponents ? t.variableComponents[c] : t.outputComponent, t.dispatchLayout.x.length === s.shape.length)).join(`
+`);
+  a.push(g), a.push(t.getUserCode());
+  const l = m(t);
+  return a.push(w(l, t)), a.join(`
+`);
+}
+function q(e, s, t) {
+  let o = e.shaderKey;
+  if (e.pixelsOpType != null)
+    return o;
+  const i = [], u = [];
+  s.forEach((a) => {
+    i.push(a.shape), u.push(a.dtype);
+  }), i.push(t.shape), u.push(t.dtype);
+  const n = s.map((a) => k(a.shape, t.shape)), r = s.map((a) => z(a.shape, t.shape)).join("_"), d = n.map((a) => a.join("_")).join(";"), p = F(e) ? "flatDispatch" : "";
+  return o += "_" + (e.workgroupSize ? e.workgroupSize.join(",") : "") + i.map((a) => a.length).join(",") + u.join(",") + e.variableNames.join(",") + d + r + p, o;
+}
+const C = `
+  struct vec5 {x: i32, y: i32, z: i32, w: i32, u: i32};
+  struct vec6 {x: i32, y: i32, z: i32, w: i32, u: i32, v: i32};
+
+  // Checks whether coordinates lie within the bounds of the shape.
+  fn coordsInBounds2D(coord : vec2<i32>, shape : vec2<i32>) -> bool {
+    return all(coord >= vec2<i32>(0)) && all(coord < shape);
+  }
+  fn coordsInBounds3D(coord : vec3<i32>, shape : vec3<i32>) -> bool {
+    return all(coord >= vec3<i32>(0)) && all(coord < shape);
+  }
+  fn coordsInBounds4D(coord : vec4<i32>, shape : vec4<i32>) -> bool {
+    return all(coord >= vec4<i32>(0)) && all(coord < shape);
+  }
+
+  fn getIndexFromCoords1D(coord : i32, shape : i32) -> i32 {
+    return coord;
+  }
+  fn getIndexFromCoords2D(coords : vec2<i32>, shape : vec2<i32>) -> i32 {
+    return dot(coords, vec2<i32>(shape.y, 1));
+  }
+  fn getIndexFromCoords3D(coords : vec3<i32>, shape : vec3<i32>) -> i32 {
+    return dot(coords, vec3<i32>(shape.y * shape.z, shape.z, 1));
+  }
+  fn getIndexFromCoords4D(coords : vec4<i32>, shape : vec4<i32>) -> i32 {
+    return dot(coords, vec4<i32>(
+        shape.y * shape.z * shape.w, shape.z * shape.w, shape.w, 1));
+  }
+  fn getIndexFromCoords5D(coords : vec5, shape : vec5) -> i32 {
+    let shapeStrides: vec5 = vec5(shape.y * shape.z * shape.w * shape.u, shape.z * shape.w * shape.u, shape.w * shape.u, shape.u, 1);
+    return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u;
+  }
+  fn getIndexFromCoords6D(coords : vec6, shape : vec6) -> i32 {
+    let shapeStrides: vec6 = vec6(shape.y * shape.z * shape.w * shape.u * shape.v, shape.z * shape.w * shape.u * shape.v, shape.w * shape.u * shape.v, shape.u * shape.v, shape.v, 1);
+    return coords.x*shapeStrides.x + coords.y*shapeStrides.y + coords.z*shapeStrides.z + coords.w*shapeStrides.w + coords.u*shapeStrides.u + coords.v*shapeStrides.v;
+  }
+
+  // NaN defination in IEEE 754-1985 is :
+  //   - sign = either 0 or 1.
+  //   - biased exponent = all 1 bits.
+  //   - fraction = anything except all 0 bits (since all 0 bits represents infinity).
+  // https://en.wikipedia.org/wiki/IEEE_754-1985#Representation_of_non-numbers
+  fn isnan(val: f32) -> bool {
+    let floatToUint: u32 = bitcast<u32>(val);
+    return (floatToUint & 0x7fffffffu) > 0x7f800000u;
+  }
+  fn isnanVec4(val : vec4<f32>) -> vec4<bool> {
+    let floatToUint: vec4<u32> = bitcast<vec4<u32>>(val);
+    return (floatToUint & vec4<u32>(0x7fffffffu)) > vec4<u32>(0x7f800000u);
+  }
+`, T = `
+  fn isinf(val: f32) -> bool {
+    return abs(val) == uniforms.INFINITY;
+  }
+`;
+function y(e, s = "") {
+  const t = e.length, o = s !== "" ? `get${s.charAt(0).toUpperCase() + s.slice(1)}CoordsFromIndex` : "getCoordsFromIndex", i = s !== "" ? `${s.charAt(0).toLowerCase() + s.slice(1)}ShapeStrides` : "outShapeStrides";
+  if (t <= 1)
+    return `fn ${o}(index : i32) -> i32 { return index; }`;
+  const u = j(e), n = v(t), r = [];
+  for (let p = 0; p < t; p++)
+    r.push(`d${p}`);
+  if (u.length === 1)
+    return `    fn ${o}(index : i32) -> vec2<i32> {
+      let d0 = index / uniforms.${i}; let d1 = index - d0 * uniforms.${i};
+      return vec2<i32>(d0, d1);
+    }`;
+  let d;
+  return d = "var index2 = index;" + u.map((p, a) => {
+    const g = `let ${r[a]} = index2 / uniforms.${i}.${I(a)}`, l = a === u.length - 1 ? `let ${r[a + 1]} = index2 - ${r[a]} * uniforms.${i}.${I(a)}` : `index2 = index2 - ${r[a]} * uniforms.${i}.${I(a)}`;
+    return `${g}; ${l};`;
+  }).join(""), `
+    fn ${o}(index : i32) -> ${n} {
+      ${d}
+      return ${n}(${r.join(",")});
+    }
+  `;
+}
+function U(e, s) {
+  const t = e.name, o = e.shape.length, i = v(o), u = "get" + t.charAt(0).toUpperCase() + t.slice(1), n = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), r = n.map((a) => `${a} : i32`).join(", ");
+  if (o < 1)
+    return `
+      fn ${u}() -> ${f(s)} {
+        return ${f(s)}(${t}[0]);
+      }
+    `;
+  const d = `uniforms.${t.charAt(0).toLowerCase() + t.slice(1)}Shape`;
+  let p = `${o}D`;
+  return o === 0 && (p = "1D"), `
+    fn ${u}(${r}) -> ${f(s)} {
+      return ${f(s)}(${t}[getIndexFromCoords${p}(${i}(${n.join(",")}),
+        ${d})${s === 1 ? "" : ` / ${s}`}]);
+    }
+   `;
+}
+function W(e, s, t, o) {
+  const i = e.name, u = i.charAt(0).toUpperCase() + i.slice(1), n = "get" + u + "ByOutput", r = e.shape.length, d = s.length, p = v(d);
+  if (z(e.shape, s) && o)
+    return `
+    fn ${n}Index(globalIndex : i32) -> ${f(t)} {
+      return ${f(t)}(${i}[globalIndex]);
+    }
+
+    fn ${n}Coords(coords : ${p}) -> ${f(t)} {
+      return ${f(t)}(${i}[${d > 1 ? "getOutputIndexFromCoords(coords)" : "coords"}${t === 1 ? "" : ` / ${t}`}]);
+    }
+    `;
+  const a = k(e.shape, s), g = d - r;
+  let l = "";
+  if (r === 0)
+    return `
+    fn ${n}Index(globalIndex : i32) -> ${f(t)}{
+      return get${u}();
+    }
+
+    fn ${n}Coords(coords : ${p}) -> ${f(t)}{
+      return get${u}();
+    }
+  `;
+  d < 2 && a.length >= 1 ? l = "coords = 0;" : l = a.map((S) => `coords.${I(S + g)} = 0;`).join(`
+`);
+  let $ = "";
+  if (d < 2 && r > 0)
+    $ = "coords";
+  else if (d > 1) {
+    const S = v(r), O = e.shape.map((M, _) => `coords.${I(_ + g)}`).join(", ");
+    $ = `${S}(${O})`;
+  } else
+    $ = "coords";
+  const h = `uniforms.${i.charAt(0).toLowerCase() + i.slice(1)}Shape`, c = `${r}D`;
+  return `
+  fn ${n}Index(globalIndex : i32) -> ${f(t)} {
+    var coords = getCoordsFromIndex(globalIndex);
+    ${l}
+    return ${f(t)}(${i}[getIndexFromCoords${c}(${$}, ${h})${t === 1 ? "" : ` / ${t}`}]);
+  }
+
+  fn ${n}Coords(coordsIn : ${p}) -> ${f(t)} {
+    var coords = coordsIn;
+    ${l}
+    return ${f(t)}(${i}[getIndexFromCoords${c}(${$}, ${h})${t === 1 ? "" : ` / ${t}`}]);
+  }
+`;
+}
+function P(e, s, t, o) {
+  let i = U(e, t);
+  return e.shape.length <= s.length && (i += W(e, s, t, o)), i;
+}
+function R(e, s) {
+  const { x: t, y: o = [], z: i = [] } = s, u = e.length, n = t.length + o.length + i.length;
+  if (n !== u)
+    return "";
+  if (t.length === u)
+    return `fn getOutputCoords() -> ${v(u)}{
+    let globalIndex = getGlobalIndex();
+    return getCoordsFromIndex(globalIndex);
+  }
+  `;
+  let r = "";
+  const d = [t, o, i];
+  for (let l = 0; l < d.length; l++) {
+    const $ = d[l];
+    if ($.length !== 0)
+      if ($.length === 1)
+        r += `let d${$[0]} = i32(globalId[${l}]);`;
+      else {
+        const h = L($, "uniforms.outShape");
+        r += `var index${l} = i32(globalId[${l}]);`;
+        for (let c = 0; c < h.length; c++)
+          r += `let d${$[c]} = index${l} / ${h[c]};`, c === h.length - 1 ? r += `let d${$[c + 1]} = index${l} - d${$[c]} * ${h[c]};` : r += `index${l} = index${l} - d${$[c]} * ${h[c]};`;
+      }
+  }
+  const p = [];
+  for (let l = 0; l < n; l++)
+    p.push(`d${l}`);
+  const a = v(n);
+  let g = `fn getOutputCoords() -> ${a} {
+  ${r}
+`;
+  return p.length === 0 ? g += `return ${a}(0); }` : g += `return ${a}(${p.join(",")}); }`, g;
+}
+function G(e) {
+  let s = "";
+  switch (e) {
+    case 0:
+    case 1:
+      s += `
+        fn getOutputIndexFromCoords(coords : i32) -> i32 {
+          return coords;
+        }
+        `;
+      break;
+    case 2:
+      s += `
+        fn getOutputIndexFromCoords(coords : vec2<i32>) -> i32 {
+          return dot(coords, vec2<i32>(uniforms.outShapeStrides, 1));
+        }
+        `;
+      break;
+    case 3:
+      s += `
+        fn getOutputIndexFromCoords(coords : vec3<i32>) -> i32 {
+          return dot(coords, vec3<i32>(uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, 1));
+        }
+        `;
+      break;
+    case 4:
+      s += `
+        fn getOutputIndexFromCoords(coords : vec4<i32>) -> i32 {
+          return dot(coords, vec4<i32>(
+            uniforms.outShapeStrides.x, uniforms.outShapeStrides.y, uniforms.outShapeStrides.z, 1));
+        }
+        `;
+      break;
+    case 5:
+      s += `
+        fn getOutputIndexFromCoords(coords : vec5) -> i32 {
+          return coords.x * uniforms.outShapeStrides.x +
+              coords.y * uniforms.outShapeStrides.y +
+              coords.z * uniforms.outShapeStrides.z +
+              coords.w * uniforms.outShapeStrides.w +
+              coords.u;
+        }
+        `;
+      break;
+    case 6:
+      s += `
+        fn getOutputIndexFromCoords(coords : vec6) -> i32 {
+          return coords.x * uniforms.outShapeStrides.x +
+              coords.y * uniforms.outShapeStrides.y +
+              coords.z * uniforms.outShapeStrides.z +
+              coords.w * uniforms.outShapeStrides.w +
+              coords.u * uniforms.outShapeStrides.u +
+              coords.v;
+        }
+        `;
+      break;
+    default:
+      A(!1, () => `Unsupported ${e}D shape`);
+      break;
+  }
+  return s;
+}
+function F(e) {
+  return e.dispatch[1] === 1 && e.dispatch[2] === 1;
+}
+function x(e, s = 1) {
+  if (e === "float32")
+    return f(s, "f32");
+  if (e === "int32" || e === "bool")
+    return f(s, "i32");
+  throw new Error(`type ${e} is not supported.`);
+}
+function V(e, s, t) {
+  const o = e.length, i = x(s, t);
+  let u = `fn setOutputAtIndex(flatIndex : i32, value : ${f(t)}) {
+      result[flatIndex] = ${i}(value);
+    }
+
+    fn setOutputAtIndexI32(flatIndex : i32, value : ${f(t, "i32")}) {
+      result[flatIndex] = ${i}(value);
+    }
+    `;
+  if (o >= 2) {
+    const n = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), r = v(o);
+    u += `
+      fn setOutputAtCoords(${n.map((d) => `${d} : i32`).join(", ")}, value : ${f(t)}) {
+        let flatIndex = getOutputIndexFromCoords(${r}(${n.join(", ")}));
+        setOutputAtIndex(flatIndex${t === 1 ? "" : ` / ${t}`}, value);
+      }
+      fn setOutputAtCoordsI32(${n.map((d) => `${d} : i32`).join(", ")}, value : ${f(t, "i32")}) {
+        let flatIndex = getOutputIndexFromCoords(${r}(${n.join(", ")}));
+        setOutputAtIndexI32(flatIndex${t === 1 ? "" : ` / ${t}`}, value);
+      }
+    `;
+  }
+  return u;
+}
+function B(e) {
+  const s = /(\w+)\s*:\s*vec(5|6)/g;
+  e = e.replace(s, (o) => "@align(16) " + o);
+  const t = /vec(5|6)\s*,\s*(\w+)/g;
+  return e = e.replace(t, (o, i, u) => `vec${i}, @align(16) ${u}`), e;
+}
+function m(e) {
+  return !(e.dispatchLayout.hasOwnProperty("y") && e.dispatchLayout.y.length !== 0 || e.dispatchLayout.hasOwnProperty("z") && e.dispatchLayout.z.length !== 0);
+}
+export {
+  b as P,
+  X as a,
+  v as b,
+  H as c,
+  I as d,
+  x as e,
+  y as f,
+  Y as g,
+  q as m,
+  f as t
+};

package/dist/webgpu_util-DhSeP4b6.js

@@ -0,0 +1,80 @@
+import { l as u } from "./index-UdZhlibC.js";
+/**
+ * @license
+ * Copyright 2019 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.
+ * =============================================================================
+ */
+const e = (r) => {
+  let t = 1;
+  for (let n = 0; n < r.length; n++)
+    t *= r[n];
+  return t;
+};
+function m(r, t, n = [1, 1, 1], a = [1, 1, 1]) {
+  const [o, i, f] = [
+    Math.ceil(e(r.x.map((c) => t[c])) / (n[0] * a[0])),
+    r.y ? Math.ceil(e(r.y.map((c) => t[c])) / (n[1] * a[1])) : 1,
+    r.z ? Math.ceil(e(r.z.map((c) => t[c])) / (n[2] * a[2])) : 1
+  ];
+  return [o, i, f];
+}
+function d(r, t, n, a = !1) {
+  const o = [8, 8, 1], i = [4, 4, 1];
+  return a || (r <= 8 && (i[1] = 1), t <= 16 && n <= 16 && (o[0] = 4)), { workgroupSize: o, elementsPerThread: i };
+}
+function p(r, t, n = !1) {
+  if (n)
+    return [8, 8, 1];
+  const a = e(r.x.map((i) => t[i])), o = e(r.y.map((i) => t[i]));
+  return a <= 4 ? [4, 16, 1] : o <= 4 ? [16, 4, 1] : [16, 16, 1];
+}
+function M(r, t, n = !1) {
+  if (n)
+    return [4, 4, 1];
+  const a = e(r.x.map((i) => t[i])), o = e(r.y.map((i) => t[i]));
+  return a <= 4 ? [1, 2, 1] : o <= 4 ? [2, 1, 1] : [2, 2, 1];
+}
+function h(r) {
+  return { x: r.map((t, n) => n) };
+}
+function x(r) {
+  if (r === "float32" || r === "int32" || r === "bool" || r === "string")
+    return 4;
+  if (r === "complex64")
+    return 8;
+  throw new Error(`Unknown dtype ${r}`);
+}
+function g() {
+  return !!(typeof globalThis < "u" && globalThis.navigator && globalThis.navigator.gpu);
+}
+function b(r, t) {
+  Array.isArray(r) || (r = [r]), r.forEach((n) => {
+    n != null && u(n.dtype !== "complex64", () => `${t} does not support complex64 tensors in the WebGPU backend.`);
+  });
+}
+var s;
+(function(r) {
+  r[r.MatMulReduceProgram = 0] = "MatMulReduceProgram", r[r.MatMulSplitKProgram = 1] = "MatMulSplitKProgram", r[r.MatMulSmallOutputSizeProgram = 2] = "MatMulSmallOutputSizeProgram", r[r.MatMulPackedProgram = 3] = "MatMulPackedProgram", r[r.MatMulMax = 4] = "MatMulMax";
+})(s || (s = {}));
+export {
+  x as G,
+  s as M,
+  d as a,
+  b,
+  m as c,
+  p as d,
+  M as e,
+  h as f,
+  g as i
+};

package/dist/{zeros-YJDE7oRb.js → zeros-DnPT2nD4.js}

@@ -1,4 +1,4 @@
-import { o as m, q as r, Z as l, E as c, _ as i, x as p, $ as u, g as x } from "./index-BAzbokzv.js";
+import { B as m, C as r, a5 as l, E as c, a6 as i, F as p, a7 as u, j as f } from "./index-UdZhlibC.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -15,13 +15,13 @@ import { o as m, q as r, Z as l, E as c, _ as i, x as p, $ as u, g as x } from "
  * limitations under the License.
  * =============================================================================
  */
-function f(a, e) {
-  const o = r(a, "real", "complex"), s = r(e, "imag", "complex");
-  l(o.shape, s.shape, `real and imag shapes, ${o.shape} and ${s.shape}, must match in call to tf.complex().`);
-  const n = { real: o, imag: s };
+function x(a, e) {
+  const s = r(a, "real", "complex"), o = r(e, "imag", "complex");
+  l(s.shape, o.shape, `real and imag shapes, ${s.shape} and ${o.shape}, must match in call to tf.complex().`);
+  const n = { real: s, imag: o };
   return c.runKernel(i, n);
 }
-const g = /* @__PURE__ */ m({ complex_: f });
+const g = /* @__PURE__ */ m({ complex_: x });
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.
@@ -40,11 +40,11 @@ const g = /* @__PURE__ */ m({ complex_: f });
  */
 function t(a, e = "float32") {
   if (p(a), e === "complex64") {
-    const s = t(a, "float32"), n = t(a, "float32");
-    return g(s, n);
+    const o = t(a, "float32"), n = t(a, "float32");
+    return g(o, n);
   }
-  const o = u(x(a), e);
-  return c.makeTensor(o, a, e);
+  const s = u(f(a), e);
+  return c.makeTensor(s, a, e);
 }
 export {
   g as c,

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@genai-fi/nanogpt",
-    "version": "0.6.3",
+    "version": "0.7.1",
     "type": "module",
     "main": "dist/main.js",
     "types": "dist/main.d.ts",
@@ -50,6 +50,7 @@
     "dependencies": {
         "@dsnp/parquetjs": "^1.8.7",
         "@tensorflow/tfjs": "^4.22.0",
+        "@tensorflow/tfjs-backend-webgpu": "^4.22.0",
         "eventemitter3": "^5.0.1",
         "jszip": "^3.10.1",
         "papaparse": "^5.5.3",