@mineproness/tensorloop 2.2.2 → 2.2.4

package/package.json

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

package/src/Softmax/index.js

@@ -113,81 +113,17 @@ export class Softmax {
     }
     async crossEntropy(probs, target) {
         const encoder = this.device.createCommandEncoder()
-
         const tokenCount = probs.size / (this.vocabSize * 4)
         const targetBuffer = await this.gpu.createBuffer(target.length * Float32Array.BYTES_PER_ELEMENT)
-        const loss = await this.gpu.createBuffer(4 * tokenCount)
-        const acc = await this.gpu.createBuffer(4 * 2)
+        if(!this.loss) this.loss = await this.gpu.createBuffer(4 * tokenCount)
+        if(!this.acc) this.acc = await this.gpu.createBuffer(4 * 1)
         await this.gpu.WriteBuffer(targetBuffer, new Uint32Array(target))
-        const bind = await this.device.createBindGroup({
-            layout: await this.lossPipline.getBindGroupLayout(0),
-            entries: [
-                {
-                    binding: 0,
-                    resource: {
-                        buffer: probs.mainBuffer
-                    }
-                },
-                {
-                    binding: 1,
-                    resource: {
-                        buffer: targetBuffer
-                    }
-                }, {
-                    binding: 2,
-                    resource: {
-                        buffer: loss
-                    }
-                }, {
-                    binding: 3,
-                    resource: {
-                        buffer: this.paramsForward
-                    }
-
-                },
-            ]
-        })
-        await this.gpu.RunPipline(encoder, this.lossPipline, bind, Math.ceil((tokenCount / 4)))
-        const outputLoss = await this.gpu.createBuffer(4)
-        const reduceBind = await this.gpu.createBindGroup(this.reduceFn, 0, [loss, outputLoss, this.paramsForward])
-        await this.gpu.RunPipline(encoder, this.reduceFn, reduceBind, 1);
-        const bind2 = await this.device.createBindGroup({
-            layout: await this.accPipline.getBindGroupLayout(0),
-            entries: [
-                {
-                    binding: 0,
-                    resource: {
-                        buffer: this.paramsForward
-                    }
-                },
-                {
-                    binding: 1,
-                    resource: {
-                        buffer: probs.mainBuffer
-                    }
-                }, {
-                    binding: 3,
-                    resource: {
-                        buffer: targetBuffer
-                    }
-                },
-                {
-                    binding: 2,
-                    resource: {
-                        buffer: acc
-                    }
-                }
-            ]
-        })
-        await this.gpu.RunPipline(encoder, this.accPipline, bind2, 1)
-        this.device.queue.submit([encoder.finish()])
-        targetBuffer.destroy()
-        loss.destroy()
-        const lossCpu = await this.gpu.readBuffer(outputLoss.size, outputLoss)
-        const accCPU = await this.gpu.readBuffer(acc.size, acc)
-        acc.destroy()
-        outputLoss.destroy()
-        return { loss: lossCpu[0], acc: accCPU[0] }
+        const bindGroup1 = await this.gpu.createBindGroup(this.lossPipline, 0, [probs.mainBuffer, targetBuffer, this.loss, this.paramsForward])
+        const bindGroup2 = await this.gpu.createBindGroup(this.accPipline, 0 , [this.paramsForward, probs.mainBuffer, targetBuffer, this.acc])
+        await this.gpu.RunPipline(encoder, this.lossPipline, bindGroup1, Math.ceil((tokenCount / 4)))
+        await this.gpu.RunPipline(encoder, this.accPipline, bindGroup2, 1)
+        await this.gpu.submitQueue(encoder)
+        return {loss: await this.gpu.readBuffer(this.loss.size, this.loss), acc: (await this.gpu.readBuffer(this.acc.size, this.acc))[0]}
     }
 
 }