@genai-fi/nanogpt 0.2.11 → 0.3.0

package/dist/ops/webgl/rope.js

@@ -0,0 +1,56 @@
+import { r as u } from "../../index-pWA4_lUh.js";
+class l {
+  variableNames = ["x", "sin", "cos"];
+  outputShape;
+  userCode;
+  // enableShapeUniforms = true;
+  customUniforms = [{ name: "pastLen", type: "int" }];
+  constructor(t, n, s, o) {
+    this.outputShape = [t, n, s, o], 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 = ${o};
+
+            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 h(e) {
+  const { x: t, sin: n, cos: s } = e.inputs, { pastLen: o } = e.attrs, a = e.backend, r = t.shape[0], d = t.shape[1], i = t.shape[2], c = t.shape[3], p = new l(r, d, i, c);
+  return a.runWebGLProgram(p, [t, n, s], "float32", [[o]]);
+}
+const f = {
+  kernelName: "Rope",
+  backendName: "webgl",
+  kernelFunc: h
+};
+u(f);

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

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

package/dist/ops/webgl/scatterSub.js

@@ -0,0 +1,27 @@
+import { r as i } from "../../index-pWA4_lUh.js";
+class u {
+  variableNames = ["labels", "softmaxProbs", "dy"];
+  outputShape;
+  userCode;
+  constructor(e, t) {
+    this.outputShape = [e, t], 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 d(o) {
+  const { logits: e, labels: t, dy: r } = o.inputs, s = o.backend, n = t.shape[0], a = e.shape[1], c = new u(n, a);
+  return s.runWebGLProgram(c, [t, e, r], "float32");
+}
+const b = {
+  kernelName: "EfficientScatterSub",
+  backendName: "webgl",
+  kernelFunc: d
+};
+i(b);