- Allow users to specify a name in node data provider json.

- Allow users to hide node data from aggregated stats table and children stats table. - Allow users to hide specific stats in node data from aggregated stats table. - In sync navigation, show a message when the mapped node cannot be found in the other side of the split pane. - Fix a bug in io highlighter. PiperOrigin-RevId: 685452581
google-ai-edge · Oct 13, 2024 · 58399da · 58399da
1 parent 7da3c4f
commit 58399da
Show file tree

Hide file tree

Showing 14 changed files with 357 additions and 110 deletions.
diff --git a/src/ui/src/components/visualizer/common/types.ts b/src/ui/src/components/visualizer/common/types.ts
@@ -236,8 +236,24 @@ export declare interface NodeInfo {
   node?: ModelNode;
 }
 
+/** Supported aggregated stats. */
+export type AggregatedStat = 'min' | 'max' | 'sum' | 'avg';
+
 /** Node data provider data for a single graph. */
 export declare interface NodeDataProviderGraphData {
+  /**
+   * The name of the node data.
+   *
+   * The node data's name (mainly used for display purposes) is determined by
+   * the following sources, in order of priority:
+   *
+   * 1) an explicitly specified name from this field, which overrides any other
+   *    source;
+   * 2) the `name` parameter in the API call, if applicable;
+   * 3) the JSON file name, if the data is loaded from a JSON file.
+   */
+  name?: string;
+
   /**
    * Node data indexed by node keys.
    *
@@ -282,7 +298,28 @@ export declare interface NodeDataProviderGraphData {
    */
   gradient?: GradientItem[];
 
-  // https://gist.github.com/Myndex/e1025706436736166561d339fd667493
+  /**
+   * Whether to hide the corresponding column in aggregated stats table
+   * (the first table).
+   *
+   * If all columns in that table are hidden, the whole table will be hidden.
+   */
+  hideInAggregatedStatsTable?: boolean;
+
+  /**
+   * Whether to hide the corresponding column in children stats table
+   * (the second table).
+   *
+   * If all columns in that table are hidden, the whole table will be hidden.
+   */
+  hideInChildrenStatsTable?: boolean;
+
+  /**
+   * The stats to hide in the aggregated stats table (the first table).
+   *
+   * The value for the hidden stat will be displayed as '-'.
+   */
+  hideAggregatedStats?: AggregatedStat[];
 }
 
 /** The top level node data provider data, indexed by graph id. */

diff --git a/src/ui/src/components/visualizer/common/utils.ts b/src/ui/src/components/visualizer/common/utils.ts
@@ -541,11 +541,14 @@ export function getOpNodeDataProviderKeyValuePairsForAttrsTable(
       type.replace(NODE_DATA_PROVIDER_SHOW_ON_NODE_TYPE_PREFIX, ''),
     );
   const runs = Object.values(curNodeDataProviderRuns).filter((run) =>
-    runNames.includes(run.runName),
+    runNames.includes(getRunName(run, {id: modelGraphId})),
   );
   for (const run of runs) {
     const value = (run.results || {})?.[modelGraphId][node.id]?.strValue || '-';
-    keyValuePairs.push({key: run.runName, value});
+    keyValuePairs.push({
+      key: getRunName(run, {id: modelGraphId}),
+      value,
+    });
   }
   return keyValuePairs;
 }
@@ -1053,3 +1056,13 @@ export function getIntersectionPoints(rect1: Rect, rect2: Rect) {
 
   return {intersection1, intersection2};
 }
+
+/** Gets the run name for the given run. */
+export function getRunName(
+  run: NodeDataProviderRunData,
+  modelGraphIdLike?: {id: string},
+): string {
+  return (
+    run.nodeDataProviderData?.[modelGraphIdLike?.id || '']?.name ?? run.runName
+  );
+}
diff --git a/src/ui/src/components/visualizer/info_panel.ts b/src/ui/src/components/visualizer/info_panel.ts
@@ -53,7 +53,12 @@ import {
   SearchMatchType,
   SearchResults,
 } from './common/types';
-import {getNamespaceLabel, isGroupNode, isOpNode} from './common/utils';
+import {
+  getNamespaceLabel,
+  getRunName,
+  isGroupNode,
+  isOpNode,
+} from './common/utils';
 import {ExpandableInfoText} from './expandable_info_text';
 import {HoverableLabel} from './hoverable_label';
 import {InfoPanelService} from './info_panel_service';
@@ -751,7 +756,7 @@ export class InfoPanel {
         nodeDataProvidersSection.items.push({
           id: run.runId,
           section: nodeDataProvidersSection,
-          label: run.runName,
+          label: getRunName(run, this.curModelGraph),
           value: strValue,
           canShowOnNode: run.done,
           showOnNode: this.curShowOnNodeDataProviderRuns[run.runId] != null,

diff --git a/src/ui/src/components/visualizer/node_data_provider_summary_panel.ng.html b/src/ui/src/components/visualizer/node_data_provider_summary_panel.ng.html
@@ -51,108 +51,116 @@
   </div>
 
   <!-- Stats table -->
-  <div class="table-container" #statsTableContainer
-      [class.collapsed]="statsTableCollapsed">
-    <div class="table-title-container">
-      <div class="title-label-container"
-          (click)="handleToggleExpandCollapseStatsTable(statsTableContainer)">
-        <mat-icon>{{statsTableTitleIcon}}</mat-icon>
-        {{statsTableTitle}}
+  @if (showStatsTable) {
+    <div class="table-container" #statsTableContainer
+        [class.collapsed]="statsTableCollapsed">
+      <div class="table-title-container">
+        <div class="title-label-container"
+            (click)="handleToggleExpandCollapseStatsTable(statsTableContainer)">
+          <mat-icon>{{statsTableTitleIcon}}</mat-icon>
+          {{statsTableTitle}}
+        </div>
       </div>
+      <table border="0" cellspacing="0" cellpadding="0" *ngIf="showResults">
+        <thead>
+          <tr>
+            <th class="stats">
+              Stat
+            </th>
+            <th *ngFor="let runItem of runItems; let i = index"
+                class="value-col"
+                [class.hidden]="runItem.hideInAggregatedStatsTable">
+              <div class="header-content">
+                <div class="index-number">{{i + 1}}</div>
+                <div class="stat-label">{{runItem.runName}}</div>
+              </div>
+            </th>
+          </tr>
+        </thead>
+        <tbody>
+          <tr *ngFor="let row of curStatRows; trackBy: trackByStat">
+            <td>{{row.stat}}</td>
+            <td *ngFor="let value of row.values; let i = index"
+                class="value-col"
+                [class.hidden]="getHideStatsTableCol(i)">
+              {{getStatValue(value)}}
+            </td>
+          </tr>
+        </tbody>
+      </table>
     </div>
-    <table border="0" cellspacing="0" cellpadding="0" *ngIf="showResults">
-      <thead>
-        <tr>
-          <th class="stats">
-            Stat
-          </th>
-          <th *ngFor="let runItem of runItems; let i = index"
-              class="value-col">
-            <div class="header-content">
-              <div class="index-number">{{i + 1}}</div>
-              <div class="stat-label">{{runItem.runName}}</div>
-            </div>
-          </th>
-        </tr>
-      </thead>
-      <tbody>
-        <tr *ngFor="let row of curStatRows; trackBy: trackByStat">
-          <td>{{row.stat}}</td>
-          <td *ngFor="let value of row.values"
-              class="value-col">
-            {{getStatValue(value)}}
-          </td>
-        </tr>
-      </tbody>
-    </table>
-  </div>
+  }
 
   <!-- Children stats table -->
-  <div class="table-container children-stats" #childrenStatsTableContainer
-      [class.collapsed]="childrenStatsTableCollapsed">
-    <div class="table-title-container">
-      <div class="title-label-container"
-          (click)="handleToggleExpandCollapseChildrenStatsTable(childrenStatsTableContainer)">
-        <mat-icon>{{childrenStatsTableTitleIcon}}</mat-icon>
-        {{childrenStatsTableTitle}}
+  @if (showChildrenStatsTable) {
+    <div class="table-container children-stats" #childrenStatsTableContainer
+        [class.collapsed]="childrenStatsTableCollapsed">
+      <div class="table-title-container">
+        <div class="title-label-container"
+            (click)="handleToggleExpandCollapseChildrenStatsTable(childrenStatsTableContainer)">
+          <mat-icon>{{childrenStatsTableTitleIcon}}</mat-icon>
+          {{childrenStatsTableTitle}}
+        </div>
+        @if (childrenStatRowsCount > tablePageSize && !childrenStatsTableCollapsed) {
+          <paginator #childrenStatsPaginator
+              [pageSize]="tablePageSize" [itemsCount]="childrenStatRowsCount"
+              (change)="handleChildrenStatsTablePaginatorChanged($event)">
+          </paginator>
+        }
       </div>
-      @if (childrenStatRowsCount > tablePageSize && !childrenStatsTableCollapsed) {
-        <paginator #childrenStatsPaginator
-            [pageSize]="tablePageSize" [itemsCount]="childrenStatRowsCount"
-            (change)="handleChildrenStatsTablePaginatorChanged($event)">
-        </paginator>
-      }
+      <ng-container *ngTemplateOutlet="nodeFilter;context:{$implicit: childrenStatsTableNodeFilter}">
+      </ng-container>
+      <table class="children-stats-results" *ngIf="showResults"
+          border="0" cellspacing="0" cellpadding="0">
+        <thead>
+          <tr>
+            <th class="index"
+                (click)="handleClickChildrenStatsHeader(-2)">
+              <div class="header-content">
+                #
+                <mat-icon *ngIf="curChildrenStatSortingColIndex === -2" class="sort">
+                  {{curChildrenStatSortingDirection === 'asc' ? 'arrow_upward' : 'arrow_downward'}}
+                </mat-icon>
+              </div>
+            </th>
+            <th class="header node-label"
+                (click)="handleClickChildrenStatsHeader(-1)">
+              <div class="header-content">
+                Node
+                <mat-icon *ngIf="curChildrenStatSortingColIndex === -1" class="sort">
+                  {{curChildrenStatSortingDirection === 'asc' ? 'arrow_upward' : 'arrow_downward'}}
+                </mat-icon>
+              </div>
+            </th>
+            <th *ngFor="let col of childrenStatsCols" class="value-col"
+                [class.hidden]="col.hideInChildrenStatsTable"
+                (click)="handleClickChildrenStatsHeader(col.colIndex)">
+              <div class="header-content">
+                <div class="index-number">{{col.runIndex + 1}}</div>
+                <div class="stat-label">{{col.label}}</div>
+                <mat-icon *ngIf="col.colIndex === curChildrenStatSortingColIndex" class="sort">
+                  {{curChildrenStatSortingDirection === 'asc' ? 'arrow_upward' : 'arrow_downward'}}
+                </mat-icon>
+              </div>
+            </th>
+          </tr>
+        </thead>
+        <tbody>
+          <tr *ngFor="let row of curPageChildrenStatRows; trackBy: trackByNodeId">
+            <td>{{row.index}}</td>
+            <td class="node-label"
+                (click)="handleClickNodeLabel(row.id)">
+              {{row.label}}
+            </td>
+            <td *ngFor="let strValue of row.colStrs; let i = index" class="value-col"
+                [class.hidden]="row.colHidden[i]">
+              {{strValue}}
+            </td>
+          </tr>
+        </tbody>
+      </table>
     </div>
-    <ng-container *ngTemplateOutlet="nodeFilter;context:{$implicit: childrenStatsTableNodeFilter}">
-    </ng-container>
-    <table class="children-stats-results" *ngIf="showResults"
-        border="0" cellspacing="0" cellpadding="0">
-      <thead>
-        <tr>
-          <th class="index"
-              (click)="handleClickChildrenStatsHeader(-2)">
-            <div class="header-content">
-              #
-              <mat-icon *ngIf="curChildrenStatSortingColIndex === -2" class="sort">
-                {{curChildrenStatSortingDirection === 'asc' ? 'arrow_upward' : 'arrow_downward'}}
-              </mat-icon>
-            </div>
-          </th>
-          <th class="header node-label"
-              (click)="handleClickChildrenStatsHeader(-1)">
-            <div class="header-content">
-              Node
-              <mat-icon *ngIf="curChildrenStatSortingColIndex === -1" class="sort">
-                {{curChildrenStatSortingDirection === 'asc' ? 'arrow_upward' : 'arrow_downward'}}
-              </mat-icon>
-            </div>
-          </th>
-          <th *ngFor="let col of childrenStatsCols" class="value-col"
-              (click)="handleClickChildrenStatsHeader(col.colIndex)">
-            <div class="header-content">
-              <div class="index-number">{{col.runIndex + 1}}</div>
-              <div class="stat-label">{{col.label}}</div>
-              <mat-icon *ngIf="col.colIndex === curChildrenStatSortingColIndex" class="sort">
-                {{curChildrenStatSortingDirection === 'asc' ? 'arrow_upward' : 'arrow_downward'}}
-              </mat-icon>
-            </div>
-          </th>
-        </tr>
-      </thead>
-      <tbody>
-        <tr *ngFor="let row of curPageChildrenStatRows; trackBy: trackByNodeId">
-          <td>{{row.index}}</td>
-          <td class="node-label"
-              (click)="handleClickNodeLabel(row.id)">
-            {{row.label}}
-          </td>
-          <td *ngFor="let strValue of row.colStrs" class="value-col">
-            {{strValue}}
-          </td>
-        </tr>
-      </tbody>
-    </table>
-  </div>
+  }
 
   <!-- Results table -->
   <div class="table-container results" #nodeDataTableContainer

diff --git a/src/ui/src/components/visualizer/node_data_provider_summary_panel.scss b/src/ui/src/components/visualizer/node_data_provider_summary_panel.scss
@@ -314,4 +314,8 @@
       font-weight: 500;
     }
   }
+
+  .hidden {
+    display: none;
+  }
 }