@genai-fi/nanogpt 0.7.0 → 0.7.1

package/dist/ops/webgpu/rope.js

@@ -1,5 +1,6 @@
-import { f as h, c as l, g as m } from "../../webgpu_util-g13LvDIv.js";
-import { f as g } from "../../index-C0dhsYom.js";
+import { g as h } from "../../webgpu_program-WaoMq-WD.js";
+import { f as l, c as m } from "../../webgpu_util-DhSeP4b6.js";
+import { f as g } from "../../index-UdZhlibC.js";
 class f {
   variableNames = ["x", "sin", "cos"];
   outputShape;
@@ -9,13 +10,13 @@ class f {
   workgroupSize = [64, 1, 1];
   size = !0;
   uniforms = "pastLen: i32";
-  constructor(t, o, s, n) {
-    this.outputShape = [t, o, s, n], this.dispatchLayout = h(this.outputShape), this.dispatch = l(this.dispatchLayout, this.outputShape, this.workgroupSize);
+  constructor(t, s, n, o) {
+    this.shaderKey = `Rope_${o}`, this.outputShape = [t, s, n, o], this.dispatchLayout = l(this.outputShape), this.dispatch = m(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     const t = this.outputShape[3];
     return `
-        ${m("index")} {
+        ${h("index")} {
             if (index < uniforms.size) {
                 let coords = getCoordsFromIndex(index); // [b, h, t, d]
                 let b = coords[0];
@@ -57,8 +58,8 @@ class f {
   }
 }
 function b(e) {
-  const { x: t, sin: o, cos: s } = e.inputs, { pastLen: n } = e.attrs, a = e.backend, i = t.shape[0], r = t.shape[1], d = t.shape[2], p = t.shape[3], u = new f(i, r, d, p), c = [{ type: "int32", data: [n] }];
-  return a.runWebGPUProgram(u, [t, o, s], "float32", c);
+  const { x: t, sin: s, cos: n } = e.inputs, { pastLen: o } = e.attrs, a = e.backend, r = t.shape[0], i = t.shape[1], d = t.shape[2], p = t.shape[3], u = new f(r, i, d, p), c = [{ type: "int32", data: [o] }];
+  return a.runWebGPUProgram(u, [t, s, n], "float32", c);
 }
 const x = {
   kernelName: "Rope",

package/dist/ops/webgpu/scatterSub.js

@@ -1,5 +1,6 @@
-import { f as c, c as u, g as d } from "../../webgpu_util-g13LvDIv.js";
-import { f as p } from "../../index-C0dhsYom.js";
+import { g as c } from "../../webgpu_program-WaoMq-WD.js";
+import { f as p, c as u } from "../../webgpu_util-DhSeP4b6.js";
+import { f as d } from "../../index-UdZhlibC.js";
 class h {
   variableNames = ["labels", "softmaxProbs", "dy"];
   outputShape;
@@ -9,11 +10,11 @@ class h {
   workgroupSize = [64, 1, 1];
   size = !0;
   constructor(t, e) {
-    this.outputShape = [t, e], this.dispatchLayout = c(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize);
+    this.outputShape = [t, e], this.dispatchLayout = p(this.outputShape), this.dispatch = u(this.dispatchLayout, this.outputShape, this.workgroupSize);
   }
   getUserCode() {
     return `
-            ${d("index")} {
+            ${c("index")} {
                 if (index < uniforms.size) {
                     let coords = getCoordsFromIndex(index); // [batch, depth]
                     let idx = i32(getLabels(coords[0]));
@@ -26,12 +27,12 @@ class h {
   }
 }
 function b(o) {
-  const { logits: t, labels: e, dy: s } = o.inputs, a = o.backend, r = e.shape[0], i = t.shape[1], n = new h(r, i);
-  return a.runWebGPUProgram(n, [e, t, s], "float32");
+  const { logits: t, labels: e, dy: s } = o.inputs, r = o.backend, a = e.shape[0], i = t.shape[1], n = new h(a, i);
+  return r.runWebGPUProgram(n, [e, t, s], "float32");
 }
 const l = {
   kernelName: "EfficientScatterSub",
   backendName: "webgpu",
   kernelFunc: b
 };
-p(l);
+d(l);

package/dist/ops/webgpu/utils/reductions.d.ts

@@ -0,0 +1,9 @@
+import { backend_util, TensorInfo } from '@tensorflow/tfjs-core';
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu_program';
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu';
+export interface ReduceWebGPUProgram extends WebGPUProgram {
+    inputShape: number[];
+}
+export declare function createReductionShader(workgroupSizeX: number, reductionOp: 'mean' | 'sum', inputSnippet: string, reducedSnippet: string, outputSnippet: string): string;
+export declare function createReduceInfo(inputs: TensorInfo[], axis: number | number[]): backend_util.ReduceInfo;
+export declare function reduce(program: ReduceWebGPUProgram, inputs: TensorInfo[], backend: WebGPUBackend): TensorInfo;

package/dist/ops/webgpu/utils/reductions.js

@@ -0,0 +1,68 @@
+import { p as l, j as d } from "../../../index-UdZhlibC.js";
+import { g as p } from "../../../webgpu_program-WaoMq-WD.js";
+import { r as f } from "../../../Reshape-Do18N3gO.js";
+import { c as x } from "../../../axis_util-5DTW2tFV.js";
+function I(e, r, t, s, u) {
+  return `
+           fn DIV_CEIL(a : u32, b : u32) -> u32 {
+            return ((a - 1u) / b + 1u);
+           }
+
+           ${`
+             var<workgroup> xBestValues : array<f32, ${e}>;
+           `}
+    
+           ${p("index")} {
+                let outputIndex = index / ${e};
+                let offset = outputIndex * uniforms.reduceSize;
+                var bestValue = 0.0;
+                let Length = uniforms.reduceSize;
+    
+                for (var k = i32(localId.x); k < Length;
+                    k = k + ${e}) {
+                    var candidate = f32(x[offset + k]);
+                    ${t}
+                    bestValue = bestValue + candidate;
+                }
+                xBestValues[localId.x] = bestValue;
+                workgroupBarrier();
+    
+                var reduceSize = min(u32(Length), ${e}u);
+                for (var currentSize = reduceSize / 2u; reduceSize > 1u;
+                    currentSize = reduceSize / 2u) {
+                    let interval = DIV_CEIL(reduceSize, 2u);
+                    if (localId.x < currentSize) {
+                        let candidate = xBestValues[localId.x + interval];
+                        bestValue = bestValue + candidate;
+                        xBestValues[localId.x] = bestValue;
+                    }
+                    reduceSize = interval;
+                    workgroupBarrier();
+                }
+
+                bestValue = xBestValues[0] ${r === "mean" ? "/ f32(uniforms.reduceSize)" : ""};
+
+                ${s}
+
+                for (var k = i32(localId.x); k < Length;
+                    k = k + ${e}) {
+                    ${u}
+                }
+           }
+         `;
+}
+function V(e, r) {
+  const t = e[0], u = l(r, t.shape), [, n] = x(t.shape, u), a = d(n), i = d(t.shape) / a;
+  return { windowSize: a, inSize: a, batchSize: i, outSize: 1 };
+}
+function $(e, r, t) {
+  const s = [], u = r[0], a = [{ type: "int32", data: [e.inputShape[1]] }], o = t.runWebGPUProgram(e, r, "float32", a);
+  s.push(o);
+  const i = f({ inputs: { x: o }, attrs: { shape: u.shape }, backend: t });
+  return s.forEach((c) => t.disposeData(c.dataId)), i;
+}
+export {
+  V as createReduceInfo,
+  I as createReductionShader,
+  $ as reduce
+};