tf.tidy

Author: alicialics

tensorboard/plugins/projector/vz_projector/knn.ts

@@ -54,21 +54,25 @@ export function findKNNGPUCosDistNorm<T>(
       .runAsyncTask(
         'Finding nearest neighbors...',
         async () => {
-          const typedArray = vector.toTypedArray(dataPoints, accessor);
-          const bigMatrix = tf.tensor(typedArray, [N, dim]);
-          const bigMatrixTransposed = tf.transpose(bigMatrix);
-          // A * A^T.
-          const cosSimilarityMatrix = tf.matMul(bigMatrix, bigMatrixTransposed);
-          bigMatrix.dispose();
-          bigMatrixTransposed.dispose();
+          const cosSimilarityMatrix = tf.tidy(() => {
+            const typedArray = vector.toTypedArray(dataPoints, accessor);
+            const bigMatrix = tf.tensor(typedArray, [N, dim]);
+            const bigMatrixTransposed = tf.transpose(bigMatrix);
+            // A * A^T.
+            return tf.matMul(bigMatrix, bigMatrixTransposed);
+          });
           // `.data()` returns flattened Float32Array of B * N dimension.
           // For matrix of
           // [ 1  2 ]
           // [ 3  4 ],
           // `.data()` returns [1, 2, 3, 4].
-          const partial = await cosSimilarityMatrix.data();
-          // Discard all tensors and free up the memory.
-          cosSimilarityMatrix.dispose();
+          let partial;
+          try {
+            partial = await cosSimilarityMatrix.data();
+          } finally {
+            // Discard all tensors and free up the memory.
+            cosSimilarityMatrix.dispose();
+          }
           for (let i = 0; i < N; i++) {
             let kMin = new KMin<NearestEntry>(k);
             for (let j = 0; j < N; j++) {