diff --git a/package.json b/package.json
index f35eba88104..3c76fd8d0db 100644
--- a/package.json
+++ b/package.json
@@ -112,10 +112,10 @@
     "@polymer/paper-toolbar": "^3.0.1",
     "@polymer/paper-tooltip": "^3.0.1",
     "@polymer/polymer": "^3.4.1",
-    "@tensorflow/tfjs": "3.4.0",
-    "@tensorflow/tfjs-backend-cpu": "3.4.0",
-    "@tensorflow/tfjs-backend-webgl": "3.4.0",
-    "@tensorflow/tfjs-core": "3.4.0",
+    "@tensorflow/tfjs": "4.22.0",
+    "@tensorflow/tfjs-backend-cpu": "4.22.0",
+    "@tensorflow/tfjs-backend-webgl": "4.22.0",
+    "@tensorflow/tfjs-core": "4.22.0",
     "@vaadin/vaadin-grid": "^20.0.2",
     "d3": "5.7.0",
     "dagre": "^0.8.5",
diff --git a/tensorboard/plugins/projector/vz_projector/knn.ts b/tensorboard/plugins/projector/vz_projector/knn.ts
index 819eb8b14ef..ad976a03697 100644
--- a/tensorboard/plugins/projector/vz_projector/knn.ts
+++ b/tensorboard/plugins/projector/vz_projector/knn.ts
@@ -71,26 +71,26 @@ export function findKNNGPUCosDistNorm<T>(
   let piece = 0;
 
   const typedArray = vector.toTypedArray(dataPoints, accessor);
-  const bigMatrix = tf.tensor(typedArray, [N, dim]);
+  const bigMatrix = tf.tensor2d(typedArray, [N, dim]);
   const bigMatrixTransposed = tf.transpose(bigMatrix);
   // 1 - A * A^T.
   const bigMatrixSquared = tf.matMul(bigMatrix, bigMatrixTransposed);
-  const cosDistMatrix = tf.sub(1, bigMatrixSquared);
+  const cosDistMatrix = tf.sub(tf.scalar(1), bigMatrixSquared);
 
-  let maybePaddedCosDistMatrix = cosDistMatrix;
+  let maybePaddedCosDistMatrix: tf.Tensor2D = cosDistMatrix;
   if (actualPieceSize * numPieces > N) {
     // Expect the input to be rank 2 (though it is not typed that way) so we
     // want to pad the first dimension so we split very evenly (all splitted
     // tensor have exactly the same dimesion).
-    const padding: Array<[number, number]> = [
+    const padding: number[][] = [
       [0, actualPieceSize * numPieces - N],
       [0, 0],
     ];
-    maybePaddedCosDistMatrix = tf.pad(cosDistMatrix, padding);
+    maybePaddedCosDistMatrix = tf.pad(cosDistMatrix, padding) as tf.Tensor2D;
   }
   const splits = tf.split(
-    maybePaddedCosDistMatrix,
-    new Array(numPieces).fill(actualPieceSize),
+    maybePaddedCosDistMatrix as tf.Tensor,
+    new Array(numPieces).fill(actualPieceSize) as number[],
     0
   );
 
diff --git a/tensorboard/webapp/third_party/tfjs.ts b/tensorboard/webapp/third_party/tfjs.ts
index b70abe02549..dba208c8095 100644
--- a/tensorboard/webapp/third_party/tfjs.ts
+++ b/tensorboard/webapp/third_party/tfjs.ts
@@ -16,3 +16,8 @@ export * from '@tensorflow/tfjs-core';
 
 import '@tensorflow/tfjs-backend-cpu';
 import '@tensorflow/tfjs-backend-webgl';
+
+// TF.js 4.x requires explicitly setting/registering a backend.
+// setBackend is re-exported from tfjs-core above; callers should invoke
+// tf.setBackend('webgl') or tf.setBackend('cpu') before running operations
+// if the default backend resolution is not sufficient.