@imaduddin/fastbrain 1.1.1 → 2.0.1

package/dist/SelfAttention/index.js

@@ -0,0 +1,186 @@
+import { join } from "node:path";
+import { Linear } from "../Linear/index.js";
+import { GPUManager } from "@imaduddin/webgpu-manager";
+export class SelfAttention {
+    device;
+    embeddingSize;
+    configs;
+    query;
+    value;
+    key;
+    v;
+    q;
+    weights;
+    k;
+    gpu;
+    params;
+    Softparams;
+    Softbackparams;
+    Transparams;
+    matrixMat;
+    transposeShader;
+    addShader;
+    softmaxShader;
+    softmaxBackShader;
+    qkT;
+    transPosedKey;
+    output;
+    qKTShader;
+    dQ;
+    dK;
+    dV;
+    dWeights;
+    dScores;
+    dInput;
+    addParams;
+    constructor(device, embeddingSize, configs) {
+        this.device = device;
+        this.embeddingSize = embeddingSize;
+        this.configs = configs;
+        this.query = new Linear(device, embeddingSize, embeddingSize, {
+            filename: [join(configs.dirname, "LINEAR_QUERY_WEIGHTS.bin"), join(configs.dirname, "LINEAR_QUERY_BIAS.bin")]
+        });
+        this.value = new Linear(device, embeddingSize, embeddingSize, {
+            filename: [join(configs.dirname, "LINEAR_VALUE_WEIGHTS.bin"), join(configs.dirname, "LINEAR_VALUE_BIAS.bin")]
+        });
+        this.key = new Linear(device, embeddingSize, embeddingSize, {
+            filename: [join(configs.dirname, "LINEAR_KEY_WEIGHTS.bin"), join(configs.dirname, "LINEAR_KEY_BIAS.bin")]
+        });
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders"
+        });
+    }
+    async softmaxBackward(weights, dWeigths, output, tokenCount, enocder) {
+        if (!this.Softbackparams) {
+            this.Softbackparams = await this.gpu?.createBuffer(new Uint32Array([tokenCount, this.embeddingSize, 0]), "uniform");
+        }
+        else {
+            await this.Softbackparams.changeData(new Uint32Array([tokenCount, this.embeddingSize, 0]));
+        }
+        if (!this.softmaxBackShader) {
+            this.softmaxBackShader = await this.gpu?.getPiplineShader("selfAttention/SB.wgsl");
+        }
+        const bind = await this.gpu?.createBindGroup(this.softmaxBackShader, [weights, dWeigths, output, this.Softbackparams], 0);
+        await this.gpu?.RunPipline(enocder, this.softmaxBackShader, bind, Math.ceil(tokenCount / 16), Math.ceil(tokenCount / 16));
+    }
+    async MatMut(a, b, output, M, N, K, encoderS, isqK) {
+        if (!this.params) {
+            this.params = await this.gpu?.createBuffer(new Uint32Array([M, N, K]), "uniform");
+            this.qKTShader = await this.gpu?.getPiplineShader("selfAttention/qK.wgsl");
+            this.matrixMat = await this.gpu?.getPiplineShader("selfAttention/matMul.wgsl");
+        }
+        else {
+            this.params.changeData(new Uint32Array([M, N, K]));
+        }
+        const encoder = encoderS ? encoderS : await this.gpu?.CreateEncoder();
+        const bindGroup = await this.gpu?.createBindGroup(isqK ? this.qKTShader : this.matrixMat, [a, b, output, this.params], 0);
+        await this.gpu?.RunPipline(encoder, isqK ? this.qKTShader : this.matrixMat, bindGroup, N, M);
+    }
+    async transpose(a, output, M, N, encoder) {
+        if (!this.Transparams) {
+            this.Transparams = await this.gpu?.createBuffer(new Uint32Array([M, N, 0]), "uniform");
+            this.matrixMat = await this.gpu?.getPiplineShader("selfAttention/matMul.wgsl");
+            this.transposeShader = await this.gpu?.getPiplineShader("selfAttention/transpose.wgsl");
+        }
+        else {
+            this.Transparams.changeData(new Uint32Array([M, N, 0]));
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.transposeShader, [a, output, this.Transparams], 0);
+        await this.gpu?.RunPipline(encoder, this.transposeShader, bindGroup, N, M);
+    }
+    async forward(input) {
+        let tokenCount = input.buffer.size / (this.embeddingSize * 4);
+        this.q = await this.query?.forward(input, undefined);
+        this.k = await this.key?.forward(input, undefined);
+        this.v = await this.value?.forward(input, undefined);
+        if (!this.softmaxShader) {
+            this.softmaxShader = await this.gpu?.getPiplineShader("softmax/forward.wgsl");
+        }
+        if (!this.qkT || this.qkT.buffer.size !== tokenCount * tokenCount * 4) {
+            if (this.qkT)
+                this.qkT.destroy();
+            this.qkT = await this.gpu?.createBuffer(new Float32Array(tokenCount * tokenCount), "storage");
+        }
+        if (!this.transPosedKey || this.transPosedKey.buffer.size !== tokenCount * this.embeddingSize * 4) {
+            if (this.transPosedKey)
+                this.transPosedKey.destroy();
+            this.transPosedKey = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.embeddingSize), "storage");
+        }
+        if (!this.weights || this.weights.buffer.size !== tokenCount * tokenCount * 4) {
+            if (this.weights)
+                this.weights.destroy();
+            this.weights = await this.gpu?.createBuffer(new Float32Array(tokenCount * tokenCount), "storage");
+        }
+        if (!this.output || this.output.buffer.size !== tokenCount * this.embeddingSize * 4) {
+            if (this.output)
+                this.output.destroy();
+            this.output = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.embeddingSize), "storage");
+        }
+        if (!this.Softparams) {
+            this.Softparams = await this.gpu?.createBuffer(new Uint32Array([tokenCount, tokenCount, 0]), "uniform");
+        }
+        else {
+            this.Softparams.changeData(new Uint32Array([tokenCount, tokenCount, 0]));
+        }
+        const encoder = await this.gpu?.CreateEncoder();
+        await this.transpose(this.k, this.transPosedKey, this.embeddingSize, tokenCount, encoder);
+        await this.MatMut(this.q, this.transPosedKey, this.qkT, tokenCount, tokenCount, this.embeddingSize, encoder, true);
+        const bindGroup = await this.gpu?.createBindGroup(this.softmaxShader, [this.qkT, this.weights, this.Softparams], 0);
+        await this.gpu?.RunPipline(encoder, this.softmaxShader, bindGroup, tokenCount);
+        await this.MatMut(this.weights, this.v, this.output, tokenCount, this.embeddingSize, tokenCount, encoder, false);
+        await this.gpu?.SumbitEncoder(encoder);
+        return this.output;
+    }
+    async Save() {
+        await this.query?.Save();
+        await this.key?.Save();
+        await this.value?.Save();
+    }
+    async add3(a, b, c, out, size, encoder) {
+        if (!this.addShader) {
+            this.addShader = await this.gpu.getPiplineShader("selfAttention/add3.wgsl");
+            this.addParams = await this.gpu.createBuffer(new Uint32Array([size]), "uniform");
+        }
+        else {
+            await this.addParams.changeData(new Uint32Array([size]));
+        }
+        const bindGroup = await this.gpu.createBindGroup(this.addShader, [a, b, c, out, this.addParams], 0);
+        await this.gpu.RunPipline(encoder, this.addShader, bindGroup, Math.ceil(size / 256));
+    }
+    async backward(outputGradient, lr) {
+        const D = this.embeddingSize;
+        const T = outputGradient.buffer.size / (D * 4);
+        if (!this.dV || this.dV.buffer.size !== T * D * 4) {
+            this.dV?.destroy();
+            this.dQ?.destroy();
+            this.dK?.destroy();
+            this.dWeights?.destroy();
+            this.dScores?.destroy();
+            this.dInput?.destroy();
+            this.dV = await this.gpu.createBuffer(new Float32Array(T * D), "storage");
+            this.dQ = await this.gpu.createBuffer(new Float32Array(T * D), "storage");
+            this.dK = await this.gpu.createBuffer(new Float32Array(T * D), "storage");
+            this.dWeights = await this.gpu.createBuffer(new Float32Array(T * T), "storage");
+            this.dScores = await this.gpu.createBuffer(new Float32Array(T * T), "storage");
+            this.dInput = await this.gpu.createBuffer(new Float32Array(T * D), "storage");
+        }
+        const encoder = await this.gpu.CreateEncoder();
+        await this.transpose(this.weights, this.qkT, T, T, encoder);
+        await this.MatMut(this.qkT, outputGradient, this.dV, T, D, T, encoder, false);
+        await this.transpose(this.v, this.transPosedKey, D, T, encoder);
+        await this.MatMut(outputGradient, this.transPosedKey, this.dWeights, T, T, D, encoder, false);
+        await this.softmaxBackward(this.weights, this.dWeights, this.dScores, T, encoder);
+        await this.MatMut(this.dScores, this.k, this.dQ, T, D, T, encoder, false);
+        await this.transpose(this.dScores, this.qkT, T, T, encoder);
+        await this.MatMut(this.qkT, this.q, this.dK, T, D, T, encoder, false);
+        const dInputQ = await this.query.backward(this.dQ, lr, encoder);
+        const dInputK = await this.key.backward(this.dK, lr, encoder);
+        const dInputV = await this.value.backward(this.dV, lr, encoder);
+        await this.add3(dInputQ, dInputK, dInputV, this.dInput, T * D, encoder);
+        await this.query.updateParameter(this.dQ, lr, encoder);
+        await this.key.updateParameter(this.dK, lr, encoder);
+        await this.value.updateParameter(this.dV, lr, encoder);
+        await this.gpu.SumbitEncoder(encoder);
+        return this.dInput;
+    }
+}

package/dist/Softmax/index.d.ts

@@ -0,0 +1,18 @@
+import { GPUManager, Buffer } from "@imaduddin/webgpu-manager";
+export declare class Softmax {
+    device: GPUDevice;
+    VocabSize: number;
+    forwardShader: GPUComputePipeline | undefined;
+    lossShader: GPUComputePipeline | undefined;
+    backwardShader: GPUComputePipeline | undefined;
+    output: Buffer | undefined;
+    loss: Buffer | undefined;
+    gpu: GPUManager | undefined;
+    forwardParams: Buffer | undefined;
+    gradOutput: Buffer | undefined;
+    targetBuffer: Buffer | undefined;
+    constructor(device: GPUDevice, VocabSize: number);
+    forward(logit: Buffer | number[], encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    backward(probs: Buffer, target: number[], encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+    crossEntropy(probs: Buffer, target: number[], encoderShared: GPUCommandEncoder | undefined): Promise<Buffer | undefined>;
+}

package/dist/Softmax/index.js

@@ -0,0 +1,96 @@
+import { GPUManager } from "@imaduddin/webgpu-manager";
+import { isBufferOrNumber } from '../fileManager.js';
+export class Softmax {
+    device;
+    VocabSize;
+    forwardShader;
+    lossShader;
+    backwardShader;
+    output;
+    loss;
+    gpu;
+    forwardParams;
+    gradOutput;
+    targetBuffer;
+    constructor(device, VocabSize) {
+        this.device = device;
+        this.VocabSize = VocabSize;
+        this.gpu = new GPUManager(device, {
+            shaderDir: './shaders/softmax'
+        });
+    }
+    async forward(logit, encoderShared) {
+        let logits = await isBufferOrNumber(logit, this.gpu);
+        const tokenCount = logits.buffer.size / (this.VocabSize * 4);
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder();
+        if (!this.forwardShader) {
+            this.forwardShader = await this.gpu?.getPiplineShader("forward.wgsl");
+        }
+        if (!this.output || this.output.buffer.size !== logits.buffer.size || this.output.isDestroy()) {
+            if (this.output)
+                this.output.destroy();
+            this.output = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.VocabSize), "storage");
+        }
+        if (!this.forwardParams) {
+            this.forwardParams = await this.gpu?.createBuffer(new Uint32Array([this.VocabSize, tokenCount]), "uniform");
+        }
+        else {
+            await this.forwardParams.changeData(new Uint32Array([this.VocabSize, tokenCount]));
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.forwardShader, [logits, this.output, this.forwardParams], 0);
+        await this.gpu?.RunPipline(encoder, this.forwardShader, bindGroup, tokenCount);
+        if (!encoderShared)
+            await this.gpu?.SumbitEncoder(encoder);
+        return this.output;
+    }
+    async backward(probs, target, encoderShared) {
+        let tokenCount = probs.buffer.size / (4 * this.VocabSize);
+        if (!this.backwardShader) {
+            this.backwardShader = await this.gpu?.getPiplineShader("backward.wgsl");
+        }
+        const encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder();
+        if (!this.gradOutput || this.gradOutput.buffer.size !== probs.buffer.size || this.gradOutput.isDestroy()) {
+            if (this.gradOutput)
+                this.gradOutput.destroy();
+            this.gradOutput = await this.gpu?.createBuffer(new Float32Array(tokenCount * this.VocabSize), "storage");
+        }
+        if (!this.targetBuffer || this.targetBuffer.buffer.size !== tokenCount * Uint32Array.BYTES_PER_ELEMENT) {
+            if (this.targetBuffer)
+                this.targetBuffer.destroy();
+            this.targetBuffer = await this.gpu?.createBuffer(new Uint32Array(target), "storage");
+        }
+        else {
+            this.targetBuffer.changeData(new Uint32Array(target));
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.backwardShader, [probs, this.gradOutput, this.targetBuffer, this.forwardParams], 0);
+        await this.gpu?.RunPipline(encoder, this.backwardShader, bindGroup, Math.ceil(this.VocabSize / 64), tokenCount);
+        if (!encoderShared)
+            await this.gpu?.SumbitEncoder(encoder);
+        return this.gradOutput;
+    }
+    async crossEntropy(probs, target, encoderShared) {
+        let encoder = encoderShared ? encoderShared : await this.gpu?.CreateEncoder();
+        let tokenCount = probs.buffer.size / (4 * this.VocabSize);
+        if (!this.targetBuffer || this.targetBuffer.buffer.size !== tokenCount * Uint32Array.BYTES_PER_ELEMENT) {
+            if (this.targetBuffer)
+                this.targetBuffer.destroy();
+            this.targetBuffer = await this.gpu?.createBuffer(new Uint32Array(target), "storage");
+        }
+        else {
+            this.targetBuffer.changeData(new Uint32Array(target));
+        }
+        if (!this.loss || this.loss.buffer.size !== tokenCount * 4) {
+            if (this.loss)
+                this.loss.destroy();
+            this.loss = await this.gpu?.createBuffer(new Float32Array(tokenCount), "storage");
+        }
+        if (!this.lossShader) {
+            this.lossShader = await this.gpu?.getPiplineShader("crossEntropy.wgsl");
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.lossShader, [probs, this.targetBuffer, this.loss, this.forwardParams], 0);
+        await this.gpu?.RunPipline(encoder, this.lossShader, bindGroup, Math.ceil(tokenCount / 16));
+        if (!encoderShared)
+            await this.gpu?.SumbitEncoder(encoder);
+        return this.loss;
+    }
+}

package/dist/Tokenizer/index.d.ts

@@ -0,0 +1,15 @@
+export declare class Tokenizer {
+    configs: {
+        vocabFile: string;
+    };
+    vocabs: string[];
+    vocabSize: number | undefined;
+    idToString: Record<number, string>;
+    stringToID: Record<string, number>;
+    constructor(configs: {
+        vocabFile: string;
+    });
+    encode(text: string): number[] | [];
+    decode(token: number[]): string;
+    trainTokenizer(text: string, file: string): void;
+}

package/dist/Tokenizer/index.js

@@ -0,0 +1,58 @@
+import fs from 'fs';
+export class Tokenizer {
+    configs;
+    vocabs = ["hi"];
+    vocabSize = 0;
+    idToString = { 12: "hi" };
+    stringToID = { "hi": 12 };
+    constructor(configs) {
+        this.configs = configs;
+        const isVocabExist = fs.existsSync(configs.vocabFile);
+        if (isVocabExist) {
+            console.log(`[LOG] Tokenizer is Loaded with ${configs.vocabFile}`);
+            this.vocabs = JSON.parse(fs.readFileSync(configs.vocabFile, "utf-8"));
+            this.vocabSize = this.vocabs?.length;
+            this.idToString = Object.fromEntries(this.vocabs?.map((e, idx) => [idx, e]));
+            this.stringToID = Object.fromEntries(this.vocabs?.map((e, idx) => [e, idx]));
+        }
+        else {
+            console.log(`[LOG] Tokenizer exited with error 1. No such File or Dir.`);
+        }
+    }
+    encode(text) {
+        const split = text.toLocaleLowerCase().split(/(\s+|[.:?,;!:'-()])/).filter(Boolean);
+        const tokens = [];
+        for (let a = 0; a < split.length; a++) {
+            const token = this.stringToID[split[a]];
+            if (token !== undefined) {
+                tokens.push(token);
+            }
+            else {
+                const unkToken = this.stringToID["<unk>"];
+                tokens.push(unkToken);
+            }
+        }
+        return tokens;
+    }
+    decode(token) {
+        let text = "";
+        for (let a = 0; a < token.length; a++) {
+            const string = this.idToString[token[a]];
+            if (string !== undefined) {
+                text += string;
+            }
+            else {
+                text += "<unk>";
+            }
+        }
+        return text;
+    }
+    trainTokenizer(text, file) {
+        let oldArray = fs.existsSync(file) ? JSON.parse(fs.readFileSync(file, 'utf-8')) : [];
+        const splitedText = text.toLocaleLowerCase().split(/\r\n|\r|\n/).filter((e) => e.trim().length > 0).map((e) => e.replaceAll("<eos>", "<eos>").split(/(\s+|[.:?,;:'!-()])/).filter(Boolean)).flat();
+        const uqiuneText = [...new Set([...splitedText, "<unk>", "_m", "_e", "<eos>", ...oldArray])];
+        console.log("Operation Completed!");
+        console.log("New Vocab Size", uqiuneText.length - oldArray.length);
+        fs.writeFileSync(file, JSON.stringify(uqiuneText));
+    }
+}

package/dist/TransformerBlock/index.d.ts

@@ -0,0 +1,21 @@
+import { FFN, RMSNorm, SelfAttention, residualConnection } from '../index.js';
+import { Buffer } from '@imaduddin/webgpu-manager';
+export declare class TransformerBlock {
+    device: GPUDevice;
+    embeddingSize: number;
+    configs: {
+        dirname: string;
+    };
+    ffn: FFN | undefined;
+    attention: SelfAttention | undefined;
+    norm: RMSNorm | undefined;
+    norm1: RMSNorm | undefined;
+    residualConnection1: residualConnection | undefined;
+    residualConnection2: residualConnection | undefined;
+    constructor(device: GPUDevice, embeddingSize: number, configs: {
+        dirname: string;
+    });
+    forward(input: Buffer): Promise<Buffer | undefined>;
+    backward(outGradient: Buffer, lr: number): Promise<Buffer | undefined>;
+    Save(): Promise<void>;
+}

package/dist/TransformerBlock/index.js

@@ -0,0 +1,43 @@
+import { FFN, RMSNorm, SelfAttention, residualConnection } from '../index.js';
+export class TransformerBlock {
+    device;
+    embeddingSize;
+    configs;
+    ffn;
+    attention;
+    norm;
+    norm1;
+    residualConnection1;
+    residualConnection2;
+    constructor(device, embeddingSize, configs) {
+        this.device = device;
+        this.embeddingSize = embeddingSize;
+        this.configs = configs;
+        this.ffn = new FFN(device, embeddingSize, configs);
+        this.norm = new RMSNorm(device, embeddingSize);
+        this.norm1 = new RMSNorm(device, embeddingSize);
+        this.attention = new SelfAttention(device, embeddingSize, configs);
+        this.residualConnection1 = new residualConnection(device);
+        this.residualConnection2 = new residualConnection(device);
+    }
+    async forward(input) {
+        const norm = await this.norm?.forward(input, undefined);
+        const att = await this.attention?.forward(norm);
+        const residual = await this.residualConnection1?.add(att, input);
+        const norm1 = await this.norm1?.forward(residual, undefined);
+        const ff = await this.ffn?.forward(norm1, undefined);
+        return await this.residualConnection2?.add(residual, ff);
+    }
+    async backward(outGradient, lr) {
+        const d_FFN = await this.ffn?.backward(outGradient, lr, undefined);
+        const d_norm1 = await this.norm1?.backward(d_FFN, undefined);
+        let d_residual2 = await this.residualConnection2?.add(outGradient, d_norm1);
+        const d_att = await this.attention?.backward(d_residual2, lr);
+        const d_norm2 = await this.norm?.backward(d_att, undefined);
+        return await this.residualConnection1?.add(outGradient, d_norm2);
+    }
+    async Save() {
+        await this.ffn?.Save();
+        await this.attention?.Save();
+    }
+}

package/dist/fileManager.d.ts

@@ -0,0 +1,6 @@
+import { Buffer as WEBB, GPUManager } from '@imaduddin/webgpu-manager';
+export declare function GetVectorsForFile(file: string, total: number, isBias: boolean | undefined, aSize: number, bSize: number): Float32Array<ArrayBuffer> | null;
+export declare function SaveFileVector(file: string, vector: Float32Array): void;
+export declare function initRandomVector(total: number): Float32Array<ArrayBuffer>;
+export declare function isBufferOrNumber(f: WEBB | number[], gpu: GPUManager): Promise<WEBB>;
+export declare function getParameter(): void;

package/dist/fileManager.js

@@ -0,0 +1,34 @@
+import fs from 'fs';
+import { join } from 'path';
+import { Buffer as WEBB } from '@imaduddin/webgpu-manager';
+process.stdout.write("\x1bc");
+let totalParameter = 0;
+export function GetVectorsForFile(file, total, isBias = false, aSize, bSize) {
+    if (!fs.existsSync(file))
+        return null;
+    totalParameter += total;
+    console.log(`Loading File ${join(file)} (${aSize}x${bSize})...`);
+    const rawBuffer = fs.readFileSync(file);
+    const vec = new Float32Array(rawBuffer.buffer, rawBuffer.byteOffset, rawBuffer.byteLength / 4);
+    if (total <= vec.length) {
+        return vec;
+    }
+    const expanded = new Float32Array(total);
+    expanded.set(vec);
+    for (let i = vec.length; i < total; i++) {
+        expanded[i] = isBias ? 0 : (Math.random() * 2 - 1) * 0.05;
+    }
+    return expanded;
+}
+export function SaveFileVector(file, vector) {
+    fs.writeFileSync(file, Buffer.from(vector.buffer, vector.byteOffset, vector.byteLength));
+}
+export function initRandomVector(total) {
+    return new Float32Array(total).map((e) => (Math.random() * 2 - 1) * 0.02);
+}
+export async function isBufferOrNumber(f, gpu) {
+    return f instanceof WEBB ? f : await gpu.createBuffer(new Float32Array(f), "storage");
+}
+export function getParameter() {
+    console.log(`Total Parameter ${(totalParameter / 1000000).toFixed(2)}M.`);
+}

package/dist/index.d.ts

@@ -0,0 +1,10 @@
+export * from './Embedding/index.js';
+export * from './Linear/index.js';
+export * from './Softmax/index.js';
+export * from './Gelu/index.js';
+export * from './FFN/index.js';
+export * from './RMSNorm/index.js';
+export * from './SelfAttention/index.js';
+export * from './Tokenizer/index.js';
+export * from './TransformerBlock/index.js';
+export * from './residualConnection/index.js';

package/dist/index.js

@@ -0,0 +1,10 @@
+export * from './Embedding/index.js';
+export * from './Linear/index.js';
+export * from './Softmax/index.js';
+export * from './Gelu/index.js';
+export * from './FFN/index.js';
+export * from './RMSNorm/index.js';
+export * from './SelfAttention/index.js';
+export * from './Tokenizer/index.js';
+export * from './TransformerBlock/index.js';
+export * from './residualConnection/index.js';

package/dist/residualConnection/index.d.ts

@@ -0,0 +1,9 @@
+import { Buffer, GPUManager } from "@imaduddin/webgpu-manager";
+export declare class residualConnection {
+    device: GPUDevice;
+    output: Buffer | undefined;
+    gpu: GPUManager | undefined;
+    addShader: GPUComputePipeline | undefined;
+    constructor(device: GPUDevice);
+    add(a: Buffer, b: Buffer): Promise<Buffer | undefined>;
+}

package/dist/residualConnection/index.js

@@ -0,0 +1,25 @@
+import { GPUManager } from "@imaduddin/webgpu-manager";
+export class residualConnection {
+    device;
+    output;
+    gpu;
+    addShader;
+    constructor(device) {
+        this.device = device;
+        this.gpu = new GPUManager(device, {
+            shaderDir: "./shaders"
+        });
+    }
+    async add(a, b) {
+        const encoder = await this.gpu?.CreateEncoder();
+        if (!this.output || this.output.buffer.size !== a.buffer.size / 4) {
+            if (this.output)
+                this.output.destroy();
+            this.output = await this.gpu?.createBuffer(new Float32Array(a.buffer.size / 4), "storage");
+        }
+        const bindGroup = await this.gpu?.createBindGroup(this.addShader, [a, b, this.output], 0);
+        await this.gpu?.RunPipline(encoder, this.addShader, bindGroup, Math.ceil((a.buffer.size / 4) / 64));
+        await this.gpu?.SumbitEncoder(encoder);
+        return this.output;
+    }
+}

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@imaduddin/fastbrain",
-  "version": "1.1.1",
+  "version": "2.0.1",
   "description": "Raw-JAVASCRIPT Based Matirx and AI LLM Package called fastBrain.",
   "keywords": [
     "ai",
@@ -12,6 +12,12 @@
     "Language",
     "Model"
   ],
+  "exports": {
+    ".": {
+      "types": "./dist/index.d.ts",
+      "import": "./dist/index.js"
+    }
+  },
   "license": "MIT",
   "author": "mineproness",
   "type": "module",
@@ -20,10 +26,21 @@
     "type": "git",
     "url": "https://github.com/imaduddin/fastBrain.git"
   },
+  "scripts": {
+    "build": "tsc",
+    "test": "node test/test1.js",
+    "test1": "node test/test2.js"
+  },
   "dependencies": {
+    "@imaduddin/webgpu-manager": "^1.0.3",
     "webgpu": "^0.6.1"
   },
   "allowScripts": {
-    "webgpu@0.6.0": true
+    "webgpu@0.6.2": true
+  },
+  "devDependencies": {
+    "@types/node": "^26.6.5",
+    "@webgpu/types": "^0.1.74",
+    "typescript": "~5.8.3"
   }
-}
+}

package/shaders/add.wgsl

@@ -0,0 +1,22 @@
+@group(0) @binding(0)
+
+var <storage, read> input : array<f32>;
+
+@group(0) @binding(1)
+
+var <storage, read> a : array<f32>;
+
+
+@group(0) @binding(1)
+
+var <storage, read> out : array<f32>;
+
+@compute @workgroup_size(64)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    if(id.x >= arrayLength(&input)){
+        return ;
+    }
+    out[id.x] = input[id.x] + a[id.x];
+
+}

package/shaders/embedding/backward.wgsl

@@ -0,0 +1,36 @@
+struct Params{
+    embeddingSize: f32,
+    vocabSize: f32,
+    lr: f32
+}
+
+@group(0) @binding(0)
+
+var <storage, read_write> vectors : array<f32>;
+
+@group(0) @binding(1)
+var <storage, read> input : array<u32>;
+
+@group(0) @binding(2)
+
+var <storage, read> dInput : array<f32>;
+
+@group(0) @binding(3)
+
+var <uniform> params : Params;
+
+@compute @workgroup_size(64, 4)
+
+fn main(@builtin(global_invocation_id) id : vec3<u32>){
+    let x = id.x;
+    let token = id.y;
+    if(x >= u32(params.embeddingSize)){
+        return;
+    }
+    if(token >= arrayLength(&input)){
+        return;
+    }
+    let tokenId = input[token];
+    vectors[tokenId * u32(params.embeddingSize) + x] -= params.lr * dInput[token * u32(params.embeddingSize) + x];
+
+}