@mineproness/tensorloop 2.1.74 → 2.1.75

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mineproness/tensorloop",
-  "version": "2.1.74",
+  "version": "2.1.75",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called Tensorloop.",
   "keywords": [
     "ai",

package/src/shaders/Norm/updateGamma.wgsl

@@ -18,21 +18,21 @@ var<workgroup> partial: array<f32, 64>;
 
 const D: u32 = %embeddingSize%u;
 const EPSILON: f32 = 1e-6;
+
 @compute @workgroup_size(64)
 fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: vec3<u32>) {
-    let dim = wid.x * 64u + local_invocation_index;
-    if (dim >= D) {
-        return;
-    }
+    let dim = wid.x * 64u + lid;
     let T = arrayLength(&sumSquares);
-    let lr = params[0];
     var sum = 0.0;
-    var token = lid;
-    while (token < T) {
-        let idx = token * D + dim;
-        let r = sqrt(sumSquares[token] / f32(D) + EPSILON);
-        sum += gradOutput[idx] * input[idx] / r;
-        token += 64u;
+    if (dim < D) {
+        var token = lid;
+
+        while (token < T) {
+            let idx = token * D + dim;
+            let r = sqrt(sumSquares[token] / f32(D) + EPSILON);
+            sum += gradOutput[idx] * input[idx] / r;
+            token += 64u;
+        }
     }
     partial[lid] = sum;
     workgroupBarrier();
@@ -47,7 +47,8 @@ fn main(@builtin(local_invocation_index) lid: u32, @builtin(workgroup_id) wid: v
         workgroupBarrier();
         stride /= 2u;
     }
-    if (lid == 0u) {
-        gamma[dim] -= lr * partial[0];
+    if (lid == 0u && wid.x * 64u < D) {
+        let dimToUpdate = wid.x * 64u;
+        gamma[dimToUpdate] -= params[0] * partial[0];
     }
 }