@imaduddin/fastbrain 1.0.8 → 1.1.0-patch

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@imaduddin/fastbrain",
-  "version": "1.0.8",
+  "version": "1.1.0-patch",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called fastBrain.",
   "keywords": [
     "ai",

package/src/PositionEmbedding/index.js

@@ -45,11 +45,16 @@ export class PositionEmbedding {
             this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
             await this.gpu.WriteBuffer(this.vecBuffer, this.vectors)
         }
+        if(!this.output) this.output = await this.gpu.ManageBuffer(input.size);
+        if(this.output.size !== input.size || this.output.isDestroyed()){
+            this.output.destroy()
+            this.output = await this.gpu.ManageBuffer(input.size)
+        }
         const size = Math.ceil((input.size / 4) / 64);
-        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input.mainBuffer])
+        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input.mainBuffer, this.output.mainBuffer])
         await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
         if(!en2) await this.gpu.submitQueue(encoder);
-        return input
+        return this.output
     }
     async backward(outputGradient, lr, en2){
         const encoder = en2 ? en2 : await this.gpu.CreateEncoder()

package/src/Softmax/index.js

@@ -124,7 +124,7 @@ export class Softmax {
         if(!encoder2) await this.gpu.submitQueue(encoder)
         return this.dLogits
     }
-    async crossEntropy(probs, target) {
+    async crossEntropy(probs, target, encoderMain) {
         const encoder = encoderMain ? encoderMain : this.device.createCommandEncoder();
         const tokenCount = probs.size / (this.vocabSize * Float32Array.BYTES_PER_ELEMENT);
         if (!Number.isInteger(tokenCount) || tokenCount <= 0) {

package/src/model/index.js

@@ -69,8 +69,8 @@ export class model {
         for (let i = this.blocks.length - 1; i >= 0; i--) {
             dInput = await this.blocks[i].backward(dInput, lr, encoder)
         };
-        await this.embedding.backward(dInput, lr)
-        await this.PosEmbedding.backward(dInput, lr)
+        await this.embedding.backward(dInput, lr, encoder)
+        await this.PosEmbedding.backward(dInput, lr, encoder)
         await this.op.submitQueue(encoder)
         return { acc, loss }
 

package/src/shaders/LB.wgsl

@@ -47,7 +47,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
 
         if (tokens > 0u) {
             let idx = weightBase + j;
-            weight[idx] -= params.lr * (grad / f32(tokens));
+            weight[idx] -= params.lr * (grad / f32(params.tokenCount));
         }
     }
 
@@ -66,6 +66,6 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid
     }
 
     if (tid == 0u && tokens > 0u) {
-        bias[outputIndex] -= params.lr * (sharedBias[0] / f32(tokens));
+        bias[outputIndex] -= params.lr * (sharedBias[0] / f32(params.tokenCount));
     }
 }

package/src/shaders/PF.wgsl

@@ -6,6 +6,11 @@ var <storage, read> posVec: array<f32>;
 var <storage, read_write> vecs : array<f32>;
 
 
+@group(0) @binding(2)
+
+var <storage, read_write> output : array<f32>;
+
+
 
 
 @compute @workgroup_size(64)
@@ -15,5 +20,5 @@ fn main(@builtin(global_invocation_id) id : vec3<u32>){
     if(x >= arrayLength(&vecs)){
         return;
     }
-    vecs[x] = vecs[x] + posVec[x];
+    output[x] = vecs[x] + posVec[x];
 }