From 811f3d0d3bc6b95d45e9859aff2345df93c34fa3 Mon Sep 17 00:00:00 2001
From: Zoey Li <31378877+alicialics@users.noreply.github.com>
Date: Sun, 26 Mar 2023 05:10:36 +0000
Subject: [PATCH] Embedding Projector: perform consistent knn calculation

---
 .../plugins/projector/vz_projector/data.ts    | 20 ++++++-------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/tensorboard/plugins/projector/vz_projector/data.ts b/tensorboard/plugins/projector/vz_projector/data.ts
index 2cbca27d7a1..6e9dde59ceb 100644
--- a/tensorboard/plugins/projector/vz_projector/data.ts
+++ b/tensorboard/plugins/projector/vz_projector/data.ts
@@ -86,7 +86,7 @@ export interface DataPoint {
     [key: string]: number;
   };
 }
-const IS_FIREFOX = navigator.userAgent.toLowerCase().indexOf('firefox') >= 0;
+
 /** Maximum sample size for each projection type. */
 export const TSNE_SAMPLE_SIZE = 10000;
 export const UMAP_SAMPLE_SIZE = 5000;
@@ -470,23 +470,15 @@ export class DataSet {
           .map((neighbors) => neighbors.slice(0, nNeighbors))
       );
     } else {
-      const knnGpuEnabled = (await util.hasWebGLSupport()) && !IS_FIREFOX;
       const numKnnNeighborsToCompute = Math.max(
         nNeighbors,
         MIN_NUM_KNN_NEIGHBORS
       );
-      const result = await (knnGpuEnabled
-        ? knn.findKNNGPUCosDistNorm(
-            data,
-            numKnnNeighborsToCompute,
-            (d) => d.vector
-          )
-        : knn.findKNN(
-            data,
-            numKnnNeighborsToCompute,
-            (d) => d.vector,
-            (a, b) => vector.cosDistNorm(a, b)
-          ));
+      const result = await knn.findKNNGPUCosDistNorm(
+        data,
+        numKnnNeighborsToCompute,
+        (d) => d.vector
+      );
       this.nearest = result;
       return Promise.resolve(
         result.map((neighbors) => neighbors.slice(0, nNeighbors))