@genai-fi/nanogpt 0.6.3 → 0.7.0

package/dist/{variable-wSS22xj5.js → variable-DXEUOwew.js}

@@ -1,4 +1,4 @@
-import { E as i } from "./index-BAzbokzv.js";
+import { E as i } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/webgpu_util-g13LvDIv.js

@@ -0,0 +1,625 @@
+import { aa as F, ab as O, ac as L, Z as N, l as _ } from "./index-C0dhsYom.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 U(e, s) {
+  if (Math.max(...e) > 5)
+    throw new Error("Cannot symbolically compute strides for rank > 6 tensor.");
+  const t = e.length, o = "xyzwuv", n = e.map((a) => `${s}.${o[a]}`), r = new Array(t - 1);
+  r[t - 2] = n[t - 1];
+  for (let a = t - 3; a >= 0; --a)
+    r[a] = `(${r[a + 1]} * ${n[a + 1]})`;
+  return r;
+}
+const H = (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 y;
+(function(e) {
+  e[e.FROM_PIXELS = 0] = "FROM_PIXELS", e[e.DRAW = 1] = "DRAW";
+})(y || (y = {}));
+const Z = (e, s, t, o, n) => {
+  const r = { dtype: o.dtype, shape: o.shape }, a = P(t, r, s), i = e.createShaderModule({ code: a, label: s.constructor.name });
+  let d = L().get("WEBGPU_PRINT_SHADER");
+  if (d !== "") {
+    d = d.toLowerCase();
+    const p = d.split(",");
+    (d === "all" || p.some((u) => s.shaderKey.toLowerCase().includes(u))) && (console.group(s.shaderKey), console.debug(a), console.groupEnd());
+  }
+  return n ? e.createComputePipelineAsync({
+    compute: { module: i, entryPoint: "_start" },
+    label: s.constructor.name,
+    layout: "auto"
+  }) : e.createComputePipeline({
+    compute: { module: i, 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 q(...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 = `
+     ${D(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 D(e) {
+  return `
+  @compute @workgroup_size(${e.workgroupSize[0]}, ${e.workgroupSize[1]}, ${e.workgroupSize[2]})
+`;
+}
+function P(e, s, t) {
+  const o = [], n = 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 {
+        ${E(t) ? "  return i32(globalId.x);" : `  return i32((workgroupId.z * numWorkgroups.x * numWorkgroups.y +
+                workgroupId.y * numWorkgroups.x + workgroupId.x) * ${n}u +
+                localIndex);
+        `}
+      }
+    `), t.pixelsOpType != null) {
+    const h = t.pixelsOpType === y.FROM_PIXELS ? `@group(0) @binding(0) var<storage, read_write> result: array<${b(s.dtype, t.outputComponent)}>;` : `@group(0) @binding(1) var<storage, read> inBuf : array<${b(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 x = k(t);
+    return [
+      C,
+      o.join(`
+`),
+      m(s.shape),
+      t.getUserCode(),
+      w(x, t)
+    ].join(`
+`);
+  }
+  let r, a, i = "struct Uniforms { NAN : f32, INFINITY : f32, ";
+  t.variableNames.forEach((h, c) => {
+    const x = v(e[c].shape.length);
+    i += `${h.charAt(0).toLowerCase() + h.slice(1)}Shape : ${x}, `, r = e[c].shape.length - 1, a = v(r), i += `${h.charAt(0).toLowerCase() + h.slice(1)}ShapeStrides: ${a}, `;
+  });
+  const d = v(s.shape.length);
+  i += `outShape : ${d}, `, r = s.shape.length - 1, a = v(r), i += `
+         outShapeStrides: ${a}, `, t.size && (i += "size : i32, "), t.uniforms && (i += t.uniforms), i += "};", i = K(i), o.push(i), 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<${b(s.dtype, t.outputComponent)}>;
+    `), t.variableNames.forEach((h, c) => {
+    o.push(`
+      @group(0) @binding(${1 + c}) var<storage, read> ${h}: array<${t.variableComponents ? b(e[c].dtype, t.variableComponents[c]) : b(e[c].dtype, t.outputComponent)}>;
+        `);
+  }), i !== "" && o.push(`
+      @group(0) @binding(${1 + t.variableNames.length}) var<uniform> uniforms: Uniforms;
+      `);
+  const p = T(s.shape, t.dispatchLayout), u = [
+    C,
+    o.join(`
+`) + W,
+    m(s.shape),
+    p,
+    V(s.shape.length)
+  ];
+  t.atomic || u.push(B(s.shape, s.dtype, t.outputComponent)), t.variableNames.forEach((h, c) => {
+    u.push(`${m(e[c].shape, h)}`);
+  });
+  const g = e.map((h, c) => G(h, s.shape, t.variableComponents ? t.variableComponents[c] : t.outputComponent, t.dispatchLayout.x.length === s.shape.length)).join(`
+`);
+  u.push(g), u.push(t.getUserCode());
+  const l = k(t);
+  return u.push(w(l, t)), u.join(`
+`);
+}
+function J(e, s, t) {
+  let o = e.shaderKey;
+  if (e.pixelsOpType != null)
+    return o;
+  const n = [], r = [];
+  s.forEach((u) => {
+    n.push(u.shape), r.push(u.dtype);
+  }), n.push(t.shape), r.push(t.dtype);
+  const a = s.map((u) => F(u.shape, t.shape)), i = s.map((u) => O(u.shape, t.shape)).join("_"), d = a.map((u) => u.join("_")).join(";"), p = E(e) ? "flatDispatch" : "";
+  return o += "_" + (e.workgroupSize ? e.workgroupSize.join(",") : "") + n.map((u) => u.length).join(",") + r.join(",") + e.variableNames.join(",") + d + i + 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);
+  }
+`, W = `
+  fn isinf(val: f32) -> bool {
+    return abs(val) == uniforms.INFINITY;
+  }
+`;
+function m(e, s = "") {
+  const t = e.length, o = s !== "" ? `get${s.charAt(0).toUpperCase() + s.slice(1)}CoordsFromIndex` : "getCoordsFromIndex", n = s !== "" ? `${s.charAt(0).toLowerCase() + s.slice(1)}ShapeStrides` : "outShapeStrides";
+  if (t <= 1)
+    return `fn ${o}(index : i32) -> i32 { return index; }`;
+  const r = N(e), a = v(t), i = [];
+  for (let p = 0; p < t; p++)
+    i.push(`d${p}`);
+  if (r.length === 1)
+    return `    fn ${o}(index : i32) -> vec2<i32> {
+      let d0 = index / uniforms.${n}; let d1 = index - d0 * uniforms.${n};
+      return vec2<i32>(d0, d1);
+    }`;
+  let d;
+  return d = "var index2 = index;" + r.map((p, u) => {
+    const g = `let ${i[u]} = index2 / uniforms.${n}.${I(u)}`, l = u === r.length - 1 ? `let ${i[u + 1]} = index2 - ${i[u]} * uniforms.${n}.${I(u)}` : `index2 = index2 - ${i[u]} * uniforms.${n}.${I(u)}`;
+    return `${g}; ${l};`;
+  }).join(""), `
+    fn ${o}(index : i32) -> ${a} {
+      ${d}
+      return ${a}(${i.join(",")});
+    }
+  `;
+}
+function M(e, s) {
+  const t = e.name, o = e.shape.length, n = v(o), r = "get" + t.charAt(0).toUpperCase() + t.slice(1), a = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), i = a.map((u) => `${u} : i32`).join(", ");
+  if (o < 1)
+    return `
+      fn ${r}() -> ${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 ${r}(${i}) -> ${f(s)} {
+      return ${f(s)}(${t}[getIndexFromCoords${p}(${n}(${a.join(",")}),
+        ${d})${s === 1 ? "" : ` / ${s}`}]);
+    }
+   `;
+}
+function R(e, s, t, o) {
+  const n = e.name, r = n.charAt(0).toUpperCase() + n.slice(1), a = "get" + r + "ByOutput", i = e.shape.length, d = s.length, p = v(d);
+  if (O(e.shape, s) && o)
+    return `
+    fn ${a}Index(globalIndex : i32) -> ${f(t)} {
+      return ${f(t)}(${n}[globalIndex]);
+    }
+
+    fn ${a}Coords(coords : ${p}) -> ${f(t)} {
+      return ${f(t)}(${n}[${d > 1 ? "getOutputIndexFromCoords(coords)" : "coords"}${t === 1 ? "" : ` / ${t}`}]);
+    }
+    `;
+  const u = F(e.shape, s), g = d - i;
+  let l = "";
+  if (i === 0)
+    return `
+    fn ${a}Index(globalIndex : i32) -> ${f(t)}{
+      return get${r}();
+    }
+
+    fn ${a}Coords(coords : ${p}) -> ${f(t)}{
+      return get${r}();
+    }
+  `;
+  d < 2 && u.length >= 1 ? l = "coords = 0;" : l = u.map((x) => `coords.${I(x + g)} = 0;`).join(`
+`);
+  let $ = "";
+  if (d < 2 && i > 0)
+    $ = "coords";
+  else if (d > 1) {
+    const x = v(i), j = e.shape.map((X, A) => `coords.${I(A + g)}`).join(", ");
+    $ = `${x}(${j})`;
+  } else
+    $ = "coords";
+  const h = `uniforms.${n.charAt(0).toLowerCase() + n.slice(1)}Shape`, c = `${i}D`;
+  return `
+  fn ${a}Index(globalIndex : i32) -> ${f(t)} {
+    var coords = getCoordsFromIndex(globalIndex);
+    ${l}
+    return ${f(t)}(${n}[getIndexFromCoords${c}(${$}, ${h})${t === 1 ? "" : ` / ${t}`}]);
+  }
+
+  fn ${a}Coords(coordsIn : ${p}) -> ${f(t)} {
+    var coords = coordsIn;
+    ${l}
+    return ${f(t)}(${n}[getIndexFromCoords${c}(${$}, ${h})${t === 1 ? "" : ` / ${t}`}]);
+  }
+`;
+}
+function G(e, s, t, o) {
+  let n = M(e, t);
+  return e.shape.length <= s.length && (n += R(e, s, t, o)), n;
+}
+function T(e, s) {
+  const { x: t, y: o = [], z: n = [] } = s, r = e.length, a = t.length + o.length + n.length;
+  if (a !== r)
+    return "";
+  if (t.length === r)
+    return `fn getOutputCoords() -> ${v(r)}{
+    let globalIndex = getGlobalIndex();
+    return getCoordsFromIndex(globalIndex);
+  }
+  `;
+  let i = "";
+  const d = [t, o, n];
+  for (let l = 0; l < d.length; l++) {
+    const $ = d[l];
+    if ($.length !== 0)
+      if ($.length === 1)
+        i += `let d${$[0]} = i32(globalId[${l}]);`;
+      else {
+        const h = U($, "uniforms.outShape");
+        i += `var index${l} = i32(globalId[${l}]);`;
+        for (let c = 0; c < h.length; c++)
+          i += `let d${$[c]} = index${l} / ${h[c]};`, c === h.length - 1 ? i += `let d${$[c + 1]} = index${l} - d${$[c]} * ${h[c]};` : i += `index${l} = index${l} - d${$[c]} * ${h[c]};`;
+      }
+  }
+  const p = [];
+  for (let l = 0; l < a; l++)
+    p.push(`d${l}`);
+  const u = v(a);
+  let g = `fn getOutputCoords() -> ${u} {
+  ${i}
+`;
+  return p.length === 0 ? g += `return ${u}(0); }` : g += `return ${u}(${p.join(",")}); }`, g;
+}
+function V(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:
+      _(!1, () => `Unsupported ${e}D shape`);
+      break;
+  }
+  return s;
+}
+function E(e) {
+  return e.dispatch[1] === 1 && e.dispatch[2] === 1;
+}
+function b(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 B(e, s, t) {
+  const o = e.length, n = b(s, t);
+  let r = `fn setOutputAtIndex(flatIndex : i32, value : ${f(t)}) {
+      result[flatIndex] = ${n}(value);
+    }
+
+    fn setOutputAtIndexI32(flatIndex : i32, value : ${f(t, "i32")}) {
+      result[flatIndex] = ${n}(value);
+    }
+    `;
+  if (o >= 2) {
+    const a = ["d0", "d1", "d2", "d3", "d4", "d5"].slice(0, o), i = v(o);
+    r += `
+      fn setOutputAtCoords(${a.map((d) => `${d} : i32`).join(", ")}, value : ${f(t)}) {
+        let flatIndex = getOutputIndexFromCoords(${i}(${a.join(", ")}));
+        setOutputAtIndex(flatIndex${t === 1 ? "" : ` / ${t}`}, value);
+      }
+      fn setOutputAtCoordsI32(${a.map((d) => `${d} : i32`).join(", ")}, value : ${f(t, "i32")}) {
+        let flatIndex = getOutputIndexFromCoords(${i}(${a.join(", ")}));
+        setOutputAtIndexI32(flatIndex${t === 1 ? "" : ` / ${t}`}, value);
+      }
+    `;
+  }
+  return r;
+}
+function K(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, n, r) => `vec${n}, @align(16) ${r}`), e;
+}
+function k(e) {
+  return !(e.dispatchLayout.hasOwnProperty("y") && e.dispatchLayout.y.length !== 0 || e.dispatchLayout.hasOwnProperty("z") && e.dispatchLayout.z.length !== 0);
+}
+/**
+ * @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 S = (e) => {
+  let s = 1;
+  for (let t = 0; t < e.length; t++)
+    s *= e[t];
+  return s;
+};
+function Q(e, s, t = [1, 1, 1], o = [1, 1, 1]) {
+  const [n, r, a] = [
+    Math.ceil(S(e.x.map((i) => s[i])) / (t[0] * o[0])),
+    e.y ? Math.ceil(S(e.y.map((i) => s[i])) / (t[1] * o[1])) : 1,
+    e.z ? Math.ceil(S(e.z.map((i) => s[i])) / (t[2] * o[2])) : 1
+  ];
+  return [n, r, a];
+}
+function ee(e, s, t, o = !1) {
+  const n = [8, 8, 1], r = [4, 4, 1];
+  return o || (e <= 8 && (r[1] = 1), s <= 16 && t <= 16 && (n[0] = 4)), { workgroupSize: n, elementsPerThread: r };
+}
+function te(e, s, t = !1) {
+  if (t)
+    return [8, 8, 1];
+  const o = S(e.x.map((r) => s[r])), n = S(e.y.map((r) => s[r]));
+  return o <= 4 ? [4, 16, 1] : n <= 4 ? [16, 4, 1] : [16, 16, 1];
+}
+function se(e, s, t = !1) {
+  if (t)
+    return [4, 4, 1];
+  const o = S(e.x.map((r) => s[r])), n = S(e.y.map((r) => s[r]));
+  return o <= 4 ? [1, 2, 1] : n <= 4 ? [2, 1, 1] : [2, 2, 1];
+}
+function oe(e) {
+  return { x: e.map((s, t) => t) };
+}
+function re(e) {
+  if (e === "float32" || e === "int32" || e === "bool" || e === "string")
+    return 4;
+  if (e === "complex64")
+    return 8;
+  throw new Error(`Unknown dtype ${e}`);
+}
+function ne() {
+  return !!(typeof globalThis < "u" && globalThis.navigator && globalThis.navigator.gpu);
+}
+function ie(e, s) {
+  Array.isArray(e) || (e = [e]), e.forEach((t) => {
+    t != null && _(t.dtype !== "complex64", () => `${s} does not support complex64 tensors in the WebGPU backend.`);
+  });
+}
+var z;
+(function(e) {
+  e[e.MatMulReduceProgram = 0] = "MatMulReduceProgram", e[e.MatMulSplitKProgram = 1] = "MatMulSplitKProgram", e[e.MatMulSmallOutputSizeProgram = 2] = "MatMulSmallOutputSizeProgram", e[e.MatMulPackedProgram = 3] = "MatMulPackedProgram", e[e.MatMulMax = 4] = "MatMulMax";
+})(z || (z = {}));
+export {
+  re as G,
+  z as M,
+  y as P,
+  Z as a,
+  ee as b,
+  Q as c,
+  H as d,
+  v as e,
+  oe as f,
+  q as g,
+  I as h,
+  ne as i,
+  ie as j,
+  te as k,
+  se as l,
+  J as m,
+  b as n,
+  m as o,
+  f as t
+};

package/dist/{zeros-YJDE7oRb.js → zeros-DCPCdFGq.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 { x as m, y as r, a0 as l, E as c, a1 as i, C as p, a2 as u, j as x } from "./index-C0dhsYom.js";
 /**
  * @license
  * Copyright 2020 Google LLC. All Rights Reserved.
@@ -16,9 +16,9 @@ import { o as m, q as r, Z as l, E as c, _ as i, x as p, $ as u, g as x } from "
  * =============================================================================
  */
 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 };
+  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 });
@@ -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(x(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.0",
     "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",