@mineproness/tensorloop 2.1.66 → 2.1.69

package/index.js

@@ -6,6 +6,9 @@ import { FFN as ffn } from "./src/FFN/index.js";
 import { Tokenizer as to } from './src/Tokenizer/index.js'
 import {SelfAttention as ST} from './src/SelfAttention/index.js'
 import {TransformerBlock as TT} from './src/TransformerBlock/index.js'
+import { Model as m } from "./src/Model/index.js";
+import {PositionEmbedding as PM } from './src/PositionEmbedding/index.js'
+import { RMSNorm as RM } from "./src/RMSNorm/index.js";
 export const Embedding = E
 export const Linear = L
 export const Softmax = S;
@@ -13,4 +16,7 @@ export const RELU = r
 export const FFN = ffn
 export const Tokenizer = to;
 export const SelfAttention = ST
-export const TransformerBlock = TT;
+export const TransformerBlock = TT;
+export const Model = m;
+export const PositionEmbedding = PM;
+export const RMSNorm = RM

package/package.json

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

package/src/Embedding/index.js

@@ -134,6 +134,7 @@ export class Embedding {
         await this.gpu.WriteBuffer(this.tokenBuffer, new Uint32Array(token));
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.embeddingSize * token.length) / 256))
         this.device.queue.submit([encoder.finish()])
-        inputGradient.destroy()
+        if(this.prevGra) this.prevGra.destroy()
+        this.prevGra = inputGradient;
     }
 }

package/src/FFN/index.js

@@ -29,10 +29,11 @@ export class FFN {
         return final
     }
     async backward(inputGradient, lr) {
+        if(this.prevOUTGR)  await this.prevOUTGR.destroy()
+        this.prevOUTGR = inputGradient;
         const logit2 = await this.linear2.backward(inputGradient, lr)
         const hidden = await this.hiddenLayer.backward(logit2)
         const final = await this.linear1.backward(hidden, lr)
-        inputGradient.destroy()
         return final
     }
     async Save(){

package/src/Linear/index.js

@@ -113,6 +113,8 @@ export class Linear {
         return outputBuffer
     }
     async backward(dOutput, learningRate) {
+        if(this.pdOutput) await this.pdOutput.destroy();
+         this.pdOutput = dOutput
         const encoder = this.device.createCommandEncoder()
         const tokenCount = dOutput.size / (this.vocabSize * 4)
         if (!this.oldParamb) {
@@ -173,7 +175,6 @@ export class Linear {
         await this.gpu.RunPipline(encoder, this.dInputBackward, bind2, Math.ceil((this.embeddingSize * tokenCount) / 256));
         await this.gpu.RunPipline(encoder, this.backwardPipline, bind, Math.ceil((this.vocabSize * tokenCount) / 256));
         await this.device.queue.submit([encoder.finish()])
-        await dOutput.destroy()
         this.input.destroy()
         return outputBuffer
 

package/src/Model/index.d.ts

@@ -0,0 +1,15 @@
+export class Model {
+    async forward(input: Number[], buffer: boolean): Promise<Float32Array | GPUBuffer>;
+    async backward(output: GPUBuffer, target: number[], lr: number): Promise<{ acc: number, loss: number }>
+    TrainTokenizer(dataset: string): void;
+    constructor(device :GPUAdapter, embeddingSize: number, configs: {
+        tokenizerName: string
+        blockSize: number,
+        dirname: string,
+        contextSize: number,
+        gpuName: string
+
+    })
+    async Save() : Promise<void>
+    async CreateContext()
+}

package/src/Model/index.js

@@ -0,0 +1,95 @@
+import { Embedding, Linear, Softmax, TransformerBlock, Tokenizer, PositionEmbedding, RMSNorm } from '../../index.js'
+import { join } from 'path'
+import fs from 'fs'
+import { OperationManager } from '../OperationManager.js'
+export class Model {
+    embedding;
+    posEmbedding;
+    linear;
+    softmax;
+    blocks = [];
+    blockSize;
+    dimSize;
+    tokenizer;
+    device;
+    configs;
+    tokenizerName;
+    op;
+    tokens;
+    rmsnorm
+    constructor(device, embeddingSize = 64, configs) {
+        this.device = device;
+        this.dimSize = embeddingSize;
+        this.blockSize = configs.blockSize;
+        this.configs = configs
+        this.tokenizer = new Tokenizer(configs.tokenizerName)
+        this.tokenizerName = configs.tokenizerName
+        this.op = new OperationManager(device)
+    }
+    async CreateContext() {
+        this.embedding = new Embedding(this.device, this.dimSize, this.tokenizer.vocabSize, {
+            save: {
+                filename: join(this.configs.dirname, "./vec.bin")
+            }
+        })
+        this.linear = new Linear(this.device, this.dimSize, this.tokenizer.vocabSize, {
+            save: {
+                filename: [join(this.configs.dirname, "weights.bin"), join(this.configs.dirname, "bias.bin")]
+            }
+        })
+        this.softmax = new Softmax(this.device, this.dimSize, this.tokenizer.vocabSize);
+        for (let a = 0; a < this.blockSize; a++) {
+            console.log(`Load the Transfomer Block ${a + 1}...`)
+            if (!fs.existsSync(join(this.configs.dirname, `block-${a}`))) fs.mkdirSync(join(this.configs.dirname, `block-${a}`));
+            this.blocks[a] = new TransformerBlock(this.device, this.dimSize, {
+                dirname: join(this.configs.dirname, `block-${a}`)
+            })
+        }
+        this.posEmbedding = new PositionEmbedding(this.device, this.dimSize, this.configs.contextSize, {
+            save:{
+                filename: join(this.configs.dirname, "posVec.bin")
+            }
+        })
+    }
+    TrainTokenizer(text = "", mergeSize = 5) {
+        this.tokenizer.trainTokenizer(text, mergeSize, this.tokenizerName)
+    }
+    async ClearInputCache() {
+        await this.linear.ClearInputCache()
+        for (let a = 0; a < this.blockSize; a++) {
+            await this.blocks[a].ClearInputCache()
+        }
+    }
+    async forward(input, readBack) {
+        if(this.configs.contextSize < input.length) throw new Error("The Input is larger then contextSize")
+        this.tokens = input;
+        await this.ClearInputCache()
+        let vec = await this.embedding.forward(input)
+        vec = await this.posEmbedding.forward(vec)
+        for (let a = 0; a < this.blockSize; a++) {
+            vec = await this.blocks[a].forward(vec)
+        }
+        let logits = await this.linear.forward(vec);
+        let probs = await this.softmax.forward(logits, readBack)
+        return probs
+    }
+    async backward(outputGradient, target, lr = 0.02) {
+        const loss = await this.softmax.crossEntropy(outputGradient, target)
+        const dLogits = await this.softmax.backward(outputGradient, target);
+        let DInput = await this.linear.backward(dLogits, lr);
+        for (let a = this.blockSize - 1; a !== 0; a--) {
+            DInput = await this.blocks[a].backward(DInput, lr);
+        }
+        this.embedding.backward(this.tokens, DInput, lr);
+        this.posEmbedding.backward(DInput, lr)
+        return loss
+    }
+    async Save() {
+        await this.embedding.Save()
+        await this.linear.Save()
+        for (let a = this.blockSize - 1; a !== 0; a--) {
+            await this.blocks[a].Save();
+        }
+        await this.posEmbedding.Save()
+    }
+}

package/src/OperationManager.js

@@ -8,9 +8,10 @@ export class OperationManager {
     device;
     matrixMatSHhaders;
     params;
-    matrixParams;
+    addShader;
     constructor(device) {
         this.device = device;
+        this.addShader = this.getShaderPipline("./src/shaders/add.wgsl")
     }
     async createBuffer(size) {
         return this.device.createBuffer({
@@ -71,18 +72,21 @@ export class OperationManager {
         pass.end()
     }
 
-    async MakeACopyBuffer(encoder, sourceBuffer) {
+    async MakeACopyBuffer(encoder2, sourceBuffer) {
+         const encoder = encoder2 ? encoder2 : await this.CreateEncoder()
         const buffer = this.device.createBuffer({
             size: sourceBuffer.size,
             usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | GPUBufferUsage.STORAGE
         })
         encoder.copyBufferToBuffer(sourceBuffer, 0, buffer, 0, sourceBuffer.size);
+        if(!encoder2) await this.submitQueue(encoder)
         return buffer
     }
     CreateEncoder() {
         return this.device.createCommandEncoder()
     }
     createBindGroup(shader, groupNum, ent) {
+        
         return this.device.createBindGroup({
             layout: shader.getBindGroupLayout(groupNum),
             entries: ent.map((e, x) => {
@@ -96,5 +100,20 @@ export class OperationManager {
     async submitQueue(encoder) {
         this.device.queue.submit([encoder.finish()])
     }
+    async add(a, b, size) {
+        if (!this.params) {
+            this.params = await this.createUniFormBuffer(4 * 4)
+
+        }
+        const encoder = await this.CreateEncoder()
+        const output = await this.createBuffer(a.size);
+        const bind = await this.createBindGroup(this.addShader, 0, [a, b, output, this.params]);
+        await this.WriteBuffer(this.params, new Uint32Array([size, 0]))
+        await this.RunPipline(encoder, this.addShader, bind, Math.ceil(size / 64));
+        await this.submitQueue(encoder)
+        // a.destroy()
+        // b.destroy()
+        return output
+    }
 
 }   

package/src/PositionEmbedding/index.d.ts

@@ -0,0 +1,10 @@
+export class PositionEmbedding{
+    constructor(device: GPUBuffer ,embedding : number, contextSize: number, config:{
+        save: {
+            filename: string
+        }
+    });
+    async forward(input: GPUBuffer) : Promise<GPUBuffer>
+    async Save() : Promise<GPUBuffer>
+    async backward(inputGradient: GPUBuffer, lr: number) : Promise<void>
+}

package/src/PositionEmbedding/index.js

@@ -0,0 +1,70 @@
+import { OperationManager } from "../OperationManager.js"
+import { getVector, SaveVector } from "../getVec.js"
+function GenerateVectors(embeddingSize, contextSize) {
+    const vec = new Float32Array(embeddingSize * contextSize).map(() => (Math.random() * 2 - 1) * 0.02)
+    return vec
+}
+
+export class PositionEmbedding {
+    vectors;
+    vecBuffer;
+    gpu;
+    embeddingSize;
+    contextSize;
+    params;
+    forwardShader;
+    backwardShader;
+    configs;
+    constructor(device, embeddingSize, contextSize, configs) {
+        this.embeddingSize = embeddingSize;
+        const total = embeddingSize * contextSize;
+        this.configs = configs;
+        this.contextSize = contextSize;
+        if (configs && configs.save) {
+            const vec = getVector(configs.save.filename, total);
+            if (vec) {
+                this.vectors = vec;
+
+            } else {
+                this.vectors = GenerateVectors(embeddingSize, contextSize)
+                SaveVector(this.configs.save.filename, this.vectors)
+            }
+
+        } else {
+            this.vectors = GenerateVectors(embeddingSize, contextSize)
+
+        }
+        this.gpu = new OperationManager(device)
+        this.forwardShader = this.gpu.getShaderPipline("./src/shaders/PF.wgsl")
+        this.backwardShader = this.gpu.getShaderPipline("./src/shaders/PB.wgsl")
+    }
+    async forward(input) {
+        const encoder = this.gpu.CreateEncoder()
+        if (!this.params) this.params = await this.gpu.createUniFormBuffer(4 * 4)
+        if (!this.vecBuffer) {
+            this.vecBuffer = await this.gpu.createBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize);
+            await this.gpu.WriteBuffer(this.vecBuffer, this.vectors)
+        }
+        const output = await this.gpu.createBuffer(input.size);
+        const size = Math.ceil((input.size / 4) / 64);
+        const bindGroup = await this.gpu.createBindGroup(this.forwardShader, 0, [this.vecBuffer, input, output])
+        if (this.prevVec) this.prevVec.destroy();
+        this.prevVec = input;
+        await this.gpu.RunPipline(encoder, this.forwardShader, bindGroup, size);
+        await this.gpu.submitQueue(encoder);
+        return output
+    }
+    async backward(outputGradient, lr){
+        const encoder = await this.gpu.CreateEncoder()
+        await this.gpu.WriteBuffer(this.params, new Float32Array([lr,0]))
+        const bindGroup = await this.gpu.createBindGroup(this.backwardShader, 0, [this.vecBuffer, outputGradient, this.params])
+        if(this.prevGradient) await this.prevGradient.destroy()
+        this.prevGradient = outputGradient;
+        await this.gpu.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((outputGradient.size / 4) / 64))
+        await this.gpu.submitQueue(encoder);
+    }
+    async Save() {
+        const vec = await this.gpu.readBuffer(Float32Array.BYTES_PER_ELEMENT * this.contextSize * this.embeddingSize, this.vecBuffer)
+        SaveVector(this.configs.save.filename, vec)
+    }
+}

package/src/RMSNorm/index.d.ts

@@ -0,0 +1,5 @@
+export class RMSNorm{
+    constructor(device : GPUAdapter, embeddingSize : number);
+    async forward(input : GPUBuffer) : Promise<GPUBuffer>
+    async backward(outputGradient : GPUBuffer) : Promise<GPUBuffer>
+}

package/src/RMSNorm/index.js

@@ -0,0 +1,45 @@
+import { OperationManager } from '../OperationManager.js'
+import fs from 'fs'
+
+export class RMSNorm {
+  embeddingSize;
+  reductionShader;
+  forwardShader;
+  backwardShader;
+  op;
+  constructor(device, embeddingSize) {
+    this.op = new OperationManager(device);
+    this.embeddingSize = embeddingSize;
+    this.reductionShader = this.op.getShaderPipline("./src/shaders/Norm/sumSquare.wgsl")
+    this.forwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsForward.wgsl")
+    this.backwardShader = this.op.getShaderPipline("./src/shaders/Norm/rmsBackward.wgsl")
+  }
+  async forward(input) {
+    const outputReduce = await this.op.createBuffer(4 * 1);
+    const output = await this.op.createBuffer(input.size);
+    const bindGroup = await this.op.createBindGroup(this.reductionShader, 0, [input, outputReduce])
+    const encoder = await this.op.CreateEncoder()
+    await this.op.RunPipline(encoder, this.reductionShader, bindGroup, Math.ceil((input.size / 4) / 64))
+    const bindGroup1 = await this.op.createBindGroup(this.forwardShader, 0, [outputReduce, input, output])
+    await this.op.RunPipline(encoder, this.forwardShader, bindGroup1, Math.ceil((input.size / 4) / 64));
+    await this.op.submitQueue(encoder);
+    if (this.rms) this.rms.destroy()
+    this.rms = outputReduce
+    if (this.input) this.input.destroy()
+    this.input = input;
+    if (this.outputPre) this.outputPre.destroy()
+    this.outputPre = output
+    return output
+  }
+  async backward(outputGradient) {
+    const output = await this.op.createBuffer(outputGradient.size);
+    const bindGroup = await this.op.createBindGroup(this.backwardShader, 0, [this.input, outputGradient, this.rms, output]);
+    const encoder = await this.op.CreateEncoder()
+    await this.op.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil((this.input.size / 4) / 64))
+    await this.op.submitQueue(encoder)
+    if (this.prevOut) this.prevOut.destroy()
+    this.prevOut = outputGradient
+    return output
+  }
+
+}

package/src/SelfAttention/index.js

@@ -102,7 +102,8 @@ export class SelfAttention {
         dK.destroy()
         dQ.destroy()
         dV.destroy()
-        dO.destroy()
+        if(this.prevdO) await this.prevdO.destroy()
+        this.prevdO = dO
         return output
         
     }

package/src/Softmax/index.js

@@ -63,6 +63,7 @@ export class Softmax {
         }
         if (cpuReadBack) {
             const outputCPU = await this.gpu.readBuffer(output.size, output);
+            logits.destroy()
             return outputCPU
         }
         return output

package/src/Tokenizer/index.js

@@ -9,7 +9,7 @@ export class Tokenizer {
     IdtoWord = {};
     WordtoID = {};
     EOSToken = 2
-    constructor(filename="") {
+    constructor(filename = "") {
         if (fs.existsSync(filename)) {
             console.log("Warning: Tokenizer is loaded with a Custom Vocab File!")
             console.log("")
@@ -22,7 +22,7 @@ export class Tokenizer {
         this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
     }
     encoder(text = "") {
-        const AddSpaces = text.toLocaleLowerCase().replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ").replaceAll("<" , " < ")
+        const AddSpaces = text.toLocaleLowerCase().replaceAll(/([,.;'{}!`~\|":><?])/g, " $1 ").replaceAll("<", " < ")
         const splitedWords = AddSpaces.split(/\s+/).filter(e => e.length > 0)
         const output = []
         for (let a = 0; a < splitedWords.length; a++) {
@@ -145,8 +145,8 @@ export class Tokenizer {
         ]
         const uniArr = [...new Set([...arrText, ...specal, ...prevVocab])];
         const output = []
-        for(let a = 0 ; a < uniArr.length; a++){
-            let text = uniArr[a].match(reg).filter((e)=> e.length > 0);
+        for (let a = 0; a < uniArr.length; a++) {
+            let text = uniArr[a].match(reg).filter((e) => e.length > 0);
             output.push(...text)
         }
         let finalOutput = [...new Set([...output])]
@@ -155,6 +155,10 @@ export class Tokenizer {
         console.log(`New VocabSize : ${finalOutput.length - prevVocab.length}`);
         console.log("_________________________________")
         console.log("Operation Completed!")
-        
+        this.vocabs = finalOutput;
+        this.vocabSize = finalOutput.length;
+        this.IdtoWord = Object.fromEntries(this.vocabs.map((e, idx) => [idx, e]))
+        this.WordtoID = Object.fromEntries(this.vocabs.map((e, idx) => [e, idx]))
+
     }
 }

package/src/TransformerBlock/index.js

@@ -1,31 +1,74 @@
-import {SelfAttention} from '../SelfAttention/index.js'
-import {FFN} from '../FFN/index.js'
-import {existsSync, mkdirSync} from 'fs'
 
-export class TransformerBlock{
+import { SelfAttention } from '../SelfAttention/index.js';
+import { FFN } from '../FFN/index.js';
+import { existsSync, mkdirSync } from 'fs';
+import { RMSNorm } from '../../index.js';
+import { OperationManager } from '../OperationManager.js';
+
+export class TransformerBlock {
     attention;
     ffn;
-    constructor(device, embeddingSize, configs){
-        if(!existsSync(configs.dirname)) mkdirSync(configs.dirname);
-        this.attention = new SelfAttention(device, embeddingSize, configs)
-        this.ffn = new FFN(device, embeddingSize, configs)
+    gpu;
+    size;
+    rmsnorm;
+    rmsnorm2;
+
+    constructor(device, embeddingSize, configs) {
+        this.size = embeddingSize;
+
+        this.rmsnorm = new RMSNorm(device, embeddingSize);
+        this.rmsnorm2 = new RMSNorm(device, embeddingSize);
+
+        if (!existsSync(configs.dirname)) {
+            mkdirSync(configs.dirname, { recursive: true });
+        }
+
+        this.attention = new SelfAttention(
+            device, embeddingSize, configs
+        );
+
+        this.ffn = new FFN(device, embeddingSize, configs);
+        this.gpu = new OperationManager(device);
     }
-    async forward(input){
-        const xAtt = await this.attention.forward(input);
-        const xFF = await this.ffn.forward(xAtt);
-        return xFF
+
+    async forward(input) {
+        const count = input.size / 4;
+        const norm1 = await this.rmsnorm.forward(input);
+        const att = await this.attention.forward(norm1);
+        const h = await this.gpu.add(input, att, count);
+        const norm2 = await this.rmsnorm2.forward(h);
+        const f = await this.ffn.forward(norm2);
+        const y = await this.gpu.add(h, f, count);
+        norm1.destroy();
+        att.destroy();
+        f.destroy();
+        return y;
     }
-    async backward(outGradient, lr){
-        const dAtt = await this.ffn.backward(outGradient, lr);
-        const dFF = await   this.attention.backward(dAtt, lr)
-        return dFF
+
+    async backward(outGradient, lr) {
+        const count = outGradient.size / 4;
+        const dNorm2 = await this.ffn.backward(outGradient, lr);
+        const dHFromFFN = await this.rmsnorm2.backward(dNorm2);
+        const dH = await this.gpu.add(outGradient, dHFromFFN, count);
+        const dNorm1 = await this.attention.backward(dH, lr);
+        const dInputFromAttention = await this.rmsnorm.backward(dNorm1);
+        const dInput = await this.gpu.add(dH, dInputFromAttention, count);
+        dNorm2.destroy();
+        dHFromFFN.destroy();
+        dH.destroy();
+        dNorm1.destroy();
+        dInputFromAttention.destroy();
+        dInputFromAttention.destroy();
+        return dInput;
     }
-    async ClearInputCache(){
-        await this.ffn.ClearInputCache()
-        await this.attention.ClearInputCache()
+
+    async ClearInputCache() {
+        await this.ffn.ClearInputCache();
+        await this.attention.ClearInputCache();
     }
-    async Save(){
-        await this.ffn.Save()
-        await this.attention.Save()
+
+    async Save() {
+        await this.ffn.Save();
+        await this.attention.Save();
     }
 }

package/src/shaders/CalculateLoss.wgsl

@@ -9,7 +9,7 @@ var <storage, read> probs : array<f32>;
 
 @group(0) @binding(1)
 
-var <storage, read> targetToken : array<f32>;
+var <storage, read> targetToken : array<u32>;
 
 @group(0) @binding(2)
 

package/src/shaders/LB.wgsl

@@ -42,16 +42,16 @@ fn main(@builtin(global_invocation_id) id: vec3<u32>) {
         let inputValue = input[token * u32(params.inputSize) + inputIndex];
         let gradient = outGradient[token * u32(params.vocabSize) + outputIndex];
         let dW = inputValue * gradient;
-        dWeights += f32(dW);
+        dWeights += dW;
     }
-    let vClap = weight[index] - params.lr * max(- 15.0, min(15.0, dWeights));
+    let vClap = weight[index] - params.lr * dWeights;
     weight[index] = vClap;
     if (inputIndex == 0u) {
         var dBias = 0.0;
         for (var token = 0u; token < u32(params.tokenCount); token++) {
-            dBias += f32(outGradient[token * u32(params.vocabSize) + outputIndex]);
+            dBias += outGradient[token * u32(params.vocabSize) + outputIndex];
         }
-        bias[outputIndex] = bias[outputIndex] - params.lr * max(- 15.0, min(15.0, dBias));
+        bias[outputIndex] = bias[outputIndex] - params.lr * dBias;
 
     }
 }

package/src/shaders/Norm/rmsBackward.wgsl

@@ -0,0 +1,35 @@
+
+@group(0) @binding(0)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read> gradOutput: array<f32>;
+
+@group(0) @binding(2)
+var<storage, read> rms: array<f32>;
+
+@group(0) @binding(3)
+var<storage, read_write> gradInput: array<f32>;
+
+var<workgroup> partial: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32) {
+    let D = arrayLength(&input);
+    partial[lid] = gradOutput[lid] * input[lid];
+    workgroupBarrier();
+    var stride = 32u;
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+        workgroupBarrier();
+        stride = stride / 2u;
+    }
+    let dot = partial[0];
+    let r = rms[0];
+    gradInput[lid] = gradOutput[lid] / r - input[lid] * dot / (f32(D) * r * r * r);
+}

package/src/shaders/Norm/rmsForward.wgsl

@@ -0,0 +1,19 @@
+@group(0) @binding(0)
+
+var <storage, read> reduce: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> vector: array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> output : array<f32>;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    let rms = sqrt(reduce[0u] / f32(arrayLength(&vector)) + 1e-6);
+    output[x] = vector[x] / rms;
+}

package/src/shaders/Norm/sumSquare.wgsl

@@ -0,0 +1,37 @@
+
+@group(0) @binding(0)
+var<storage, read> input: array<f32>;
+
+@group(0) @binding(1)
+var<storage, read_write> out: array<f32>;
+
+var<workgroup> partial: array<f32, 64>;
+
+@compute @workgroup_size(64)
+fn main(@builtin(local_invocation_index) lid: u32) {
+    var sum = 0.0;
+    var i = lid;
+    while (i < arrayLength(&input)) {
+        sum += input[i] * input[i];
+        i += 64u;
+    }
+    partial[lid] = sum;
+    workgroupBarrier();
+    var stride = 32u;
+
+    loop {
+        if (stride == 0u) {
+            break;
+        }
+
+        if (lid < stride) {
+            partial[lid] += partial[lid + stride];
+        }
+
+        workgroupBarrier();
+        stride = stride / 2u;
+    }
+    if (lid == 0u) {
+        out[0] = partial[0];
+    }
+}

package/src/shaders/PB.wgsl

@@ -0,0 +1,25 @@
+struct Params{
+    lr : f32
+}
+
+@group(0) @binding(0)
+var <storage, read_write> posVec: array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> outputGradient : array<f32>;
+
+
+@group(0) @binding(2)
+
+var <uniform> params: Params;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    if(x >= arrayLength(&outputGradient)){
+        return;
+    }
+    posVec[x] -= params.lr * outputGradient[x];
+}

package/src/shaders/PF.wgsl

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

package/src/shaders/add.wgsl

@@ -0,0 +1,30 @@
+struct Params{
+  size: u32
+}
+
+@group(0) @binding(0)
+
+var <storage, read> a : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> b : array<f32>;
+
+@group(0) @binding(2)
+
+var <storage, read_write> output : array<f32>;
+@group(0) @binding(3)
+
+var <uniform> params : Params;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let index = id.x;
+    if(index >= params.size){
+        return;
+    }
+    let mainOut =  a[index] + b[index];
+    output[index] = mainOut; 
+}
+