From 76fb16d3386618521061ca248757809014bbadcc Mon Sep 17 00:00:00 2001
From: Rob Currie <robert.currie@gmail.com>
Date: Thu, 19 Dec 2024 10:18:46 -0700
Subject: [PATCH] Fix attention accumulator un-initialized logic, add comment
 on why onnx webgpu won't work for now

---
 README.md | 4 ++++
 worker.js | 9 ++++++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index ac30a5e..29db37a 100644
--- a/README.md
+++ b/README.md
@@ -38,6 +38,10 @@ make serve
 
 Processing a test sample with 2638 cells took 67 seconds in the browser vs. 34 seconds in python on the same machine.
 
+# Leveraging a GPU
+
+ONNX Web Runtime does have support for GPUs, but unfortunately they don't support all operators yet. Specifically TopK is not [supported](https://github.com/microsoft/onnxruntime/blob/main/js/web/docs/webgpu-operators.md)
+
 # References
 
 [Open Neural Network Exchange (ONNX)](https://onnx.ai/)
diff --git a/worker.js b/worker.js
index a6ee4fb..40caa00 100644
--- a/worker.js
+++ b/worker.js
@@ -1,5 +1,4 @@
 self.importScripts(
-  "https://cdnjs.cloudflare.com/ajax/libs/onnxruntime-web/1.20.1/ort.min.js",
   "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js",
   "https://cdn.jsdelivr.net/npm/h5wasm@0.7.8/dist/iife/h5wasm.min.js",
   "https://cdn.jsdelivr.net/npm/umap-js@1.4.0/lib/umap-js.min.js"
@@ -22,7 +21,7 @@ self.addEventListener("message", async function (event) {
       genes: self.model.genes,
     });
   } else if (type === "resetAttentionAccumulator") {
-    attentionAccumulator = new Float32Array(self.model.genes.length);
+    attentionAccumulator = null;
     self.postMessage({ type: "attentionAccumulatorReset" });
   }
 });
@@ -251,6 +250,10 @@ async function predict(event) {
 
       encodings.push(output.encoding.cpuData);
 
+      if (!attentionAccumulator) {
+        attentionAccumulator = new Float32Array(genes.length);
+      }
+
       for (let i = 0; i < attentionAccumulator.length; i++) {
         attentionAccumulator[i] += output.attention.cpuData[i];
       }
@@ -300,4 +303,4 @@ async function predict(event) {
     FS.unmount("/work");
     self.postMessage({ type: "error", error: error.message });
   }
-};
+}