Fix projector legend point coloring (#6021)

Fixes a regression where sprites are shown as black squares when first
loaded and when any sprite is deselected. This ultimately stemmed from a
null check (which colored the legend point white when null) being passed
over when `strictNullChecks` was introduced.

This allows `legendPointColorer` to take a null value. If `colorOption`
does not exist, it doesn't make sense to still go over the full dataSet
color assignment. Instead, return early to get `noSelectionColor`.

Googlers, see b/255947072.

Author: groszewn

tensorboard/plugins/projector/vz_projector/projectorScatterPlotAdapter.ts

@@ -88,7 +88,7 @@ export class ProjectorScatterPlotAdapter {
   private neighborsOfFirstSelectedPoint: knn.NearestEntry[];
   private renderLabelsIn3D: boolean = false;
   private labelPointAccessor: string | null;
-  private legendPointColorer: (ds: DataSet, index: number) => string;
+  private legendPointColorer: ((ds: DataSet, index: number) => string) | null;
   private distanceMetric: DistanceFunction;
   private spriteVisualizer: ScatterPlotVisualizerSprites;
   private labels3DVisualizer: ScatterPlotVisualizer3DLabels;
@@ -176,7 +176,7 @@ export class ProjectorScatterPlotAdapter {
     this.scatterPlot.render();
   }
   setLegendPointColorer(
-    legendPointColorer: (ds: DataSet, index: number) => string
+    legendPointColorer: ((ds: DataSet, index: number) => string) | null
   ) {
     this.legendPointColorer = legendPointColorer;
   }
@@ -477,7 +477,7 @@ export class ProjectorScatterPlotAdapter {
   }
   generateLineSegmentColorMap(
     ds: DataSet,
-    legendPointColorer: (ds: DataSet, index: number) => string
+    legendPointColorer: ((ds: DataSet, index: number) => string) | null
   ): {
     [polylineIndex: number]: Float32Array;
   } {
@@ -566,7 +566,7 @@ export class ProjectorScatterPlotAdapter {
   }
   generatePointColorArray(
     ds: DataSet,
-    legendPointColorer: (ds: DataSet, index: number) => string,
+    legendPointColorer: ((ds: DataSet, index: number) => string) | null,
     distFunc: DistanceFunction,
     selectedPointIndices: number[],
     neighborsOfFirstPoint: knn.NearestEntry[],

tensorboard/plugins/projector/vz_projector/vz-projector.ts

@@ -465,7 +465,10 @@ class Projector
   }
   private getLegendPointColorer(
     colorOption: ColorOption
-  ): (ds: DataSet, index: number) => string {
+  ): ((ds: DataSet, index: number) => string) | null {
+    if (colorOption == null || colorOption.map == null) {
+      return null;
+    }
     const colorer = (ds: DataSet, i: number) => {
       let value = ds.points[i].metadata[this.selectedColorOption.name];
       if (value == null) {