Add support for a special "node ids" node attribute type.

Each entry will allow users to click to reval/select the corresponding node.

PiperOrigin-RevId: 733463893

Author: Google AI Edge

src/ui/src/components/visualizer/common/input_graph.ts

@@ -23,6 +23,7 @@ import {
   IncomingEdge,
   KeyValueList,
   MetadataItem,
+  NodeAttributeList,
 } from './types';
 
 /** A collection of graphs. This is the input to the visualizer. */
@@ -136,7 +137,7 @@ export declare interface GraphNode {
   subgraphIds?: string[];
 
   /** The attributes of the node.  */
-  attrs?: KeyValueList;
+  attrs?: NodeAttributeList;
 
   /** A list of incoming edges. */
   incomingEdges?: IncomingEdge[];

src/ui/src/components/visualizer/common/model_graph.ts

@@ -22,6 +22,7 @@ import {
   GroupNodeAttributes,
   IncomingEdge,
   KeyValuePairs,
+  NodeAttributePairs,
   OutgoingEdge,
   Point,
 } from './types';
@@ -198,7 +199,7 @@ export declare interface OpNode extends ModelNodeBase {
   outgoingEdges?: OutgoingEdge[];
 
   /** The attributes of the node. */
-  attrs?: KeyValuePairs;
+  attrs?: NodeAttributePairs;
 
   /**
    * Metadata for inputs, indexed by input ids. Each input can have multiple

src/ui/src/components/visualizer/common/types.ts

@@ -31,6 +31,39 @@ export declare interface KeyValue {
 /** A type for a list of key-value pairs. */
 export type KeyValueList = KeyValue[];
 
+/** A list of node attributes. */
+export type NodeAttributeList = NodeAttribute[];
+
+/** Node attributes as a record. */
+export type NodeAttributePairs = Record<string, NodeAttributeValue>;
+
+/** A type for a single node attribute. */
+export declare interface NodeAttribute {
+  key: string;
+  value: NodeAttributeValue;
+}
+
+/** A single node attribute value. */
+export type NodeAttributeValue = string | SpecialNodeAttributeValue;
+
+/** non-string node attribute value. */
+export type SpecialNodeAttributeValue = NodeIdsNodeAttributeValue;
+
+/** Node attribute value types. */
+export enum NodeAttributeValueType {
+  NODE_IDS = 'node_ids',
+}
+
+/**
+ * A "node ids" node attribute value.
+ *
+ * Clicking on a node id will jump to the corresponding node in the graph.
+ */
+export declare interface NodeIdsNodeAttributeValue {
+  type: NodeAttributeValueType.NODE_IDS;
+  nodeIds: string[];
+}
+
 /** An item in input/output metadata. */
 export interface MetadataItem {
   id: string;

src/ui/src/components/visualizer/common/utils.ts

@@ -39,6 +39,7 @@ import {
   FieldLabel,
   KeyValueList,
   KeyValuePairs,
+  NodeAttributeValueType,
   NodeDataProviderResultProcessedData,
   NodeDataProviderRunData,
   NodeDataProviderValueInfo,
@@ -401,24 +402,26 @@ export function getOpNodeAttrsKeyValuePairsForAttrsTable(
   for (const attrId of Object.keys(attrs)) {
     const key = attrId;
     const value = attrs[attrId];
-    const matchTargets = [`${key}:${value}`, `${key}=${value}`];
-    if (
-      filterRegex.trim() === '' ||
-      matchTargets.some((matchTarget) => regex.test(matchTarget))
-    ) {
-      // Remove new line chars and spaces.
-      let processedValue = value;
-      if (key === TENSOR_VALUES_KEY) {
-        // For __value attribute, remove all white space chars.
-        processedValue = value.replace(/\s/gm, '');
-      } else {
-        // For other attributes, only remove newline chars.
-        processedValue = value.replace(/(\r\n|\n|\r)/gm, ' ');
+    if (typeof value === 'string') {
+      const matchTargets = [`${key}:${value}`, `${key}=${value}`];
+      if (
+        filterRegex.trim() === '' ||
+        matchTargets.some((matchTarget) => regex.test(matchTarget))
+      ) {
+        // Remove new line chars and spaces.
+        let processedValue = value;
+        if (key === TENSOR_VALUES_KEY) {
+          // For __value attribute, remove all white space chars.
+          processedValue = value.replace(/\s/gm, '');
+        } else {
+          // For other attributes, only remove newline chars.
+          processedValue = value.replace(/(\r\n|\n|\r)/gm, ' ');
+        }
+        keyValuePairs.push({
+          key,
+          value: processedValue,
+        });
       }
-      keyValuePairs.push({
-        key,
-        value: processedValue,
-      });
     }
   }
   return keyValuePairs;
@@ -784,7 +787,19 @@ export function getAttributesFromNode(
 ): KeyValuePairs {
   let attrs: KeyValuePairs = {};
   if (isOpNode(node)) {
-    attrs = {...(node.attrs || {})};
+    for (const [key, value] of Object.entries(node.attrs || {})) {
+      if (typeof value === 'string') {
+        attrs[key] = value;
+      } else {
+        switch (value.type) {
+          case NodeAttributeValueType.NODE_IDS:
+            attrs[key] = value.nodeIds.join(',');
+            break;
+          default:
+            break;
+        }
+      }
+    }
     // Add id to attribute.
     attrs['id'] = node.id;
   } else if (isGroupNode(node)) {

src/ui/src/components/visualizer/info_panel.ng.html

@@ -34,11 +34,30 @@
               <tr [class.search-match]="isSearchMatchedAttrId(item.label)">
                 <td class="key"><hoverable-label [label]="item.label"></hoverable-label></td>
                 <td class="value">
-                  <expandable-info-text
-                      [text]="item.value" [type]="item.label"
-                      [bgColor]="item.bgColor || 'transparent'"
-                      [textColor]="item.textColor || 'black'">
-                  </expandable-info-text>
+                  @if (item.specialValue) {
+                    @switch (item.specialValue.type) {
+                      @case (NodeAttributeValueType.NODE_IDS) {
+                        @for (nodeId of item.specialValue.nodeIds; track nodeId) {
+                          <div class="node-id-attribute"
+                              (click)="handleLocateNode(nodeId, $event)">
+                            {{nodeId}}
+                            <div class="locator-icon-container"
+                                [matTooltip]="locatorTooltip"
+                                matTooltipClass="multiline-tooltip-left"
+                                matTooltipPosition="right">
+                              <mat-icon class="locator-icon">my_location</mat-icon>
+                            </div>
+                          </div>
+                        }
+                      }
+                    }
+                  } @else {
+                    <expandable-info-text
+                        [text]="item.value" [type]="item.label"
+                        [bgColor]="item.bgColor || 'transparent'"
+                        [textColor]="item.textColor || 'black'">
+                    </expandable-info-text>
+                  }
                 </td>
               </tr>
             }

src/ui/src/components/visualizer/info_panel.scss

@@ -125,6 +125,11 @@
         margin-right: 0;
       }
     }
+
+    .node-id-attribute {
+      display: flex;
+      cursor: pointer;
+    }
   }
 
   .section:not(:first-child) {

src/ui/src/components/visualizer/info_panel.ts

@@ -47,13 +47,16 @@ import {
   KeyValue,
   KeyValueList,
   KeyValuePairs,
+  NodeAttributeValueType,
   NodeDataProviderRunInfo,
+  NodeIdsNodeAttributeValue,
   OutgoingEdge,
   SearchMatchAttr,
   SearchMatchInputMetadata,
   SearchMatchOutputMetadata,
   SearchMatchType,
   SearchResults,
+  SpecialNodeAttributeValue,
 } from './common/types';
 import {
   getNamespaceLabel,
@@ -101,6 +104,7 @@ interface InfoItem {
   bgColor?: string;
   textColor?: string;
   loading?: boolean;
+  specialValue?: SpecialNodeAttributeValue;
 }
 
 interface OutputItem {
@@ -155,6 +159,8 @@ export class InfoPanel {
   @ViewChildren('inputValueContent')
   inputValueContents = new QueryList<ElementRef<HTMLElement>>();
 
+  readonly NodeAttributeValueType = NodeAttributeValueType;
+
   private curModelGraph?: ModelGraph;
   private curSelectedNode?: ModelNode;
   private readonly curShowOnOpNodeInfoIds = new Set<string>();
@@ -769,12 +775,17 @@ export class InfoPanel {
         if (key.startsWith('__')) {
           continue;
         }
+        const value = attrs[key];
+        const strValue = typeof value === 'string' ? value : '';
+        const specialValue: SpecialNodeAttributeValue | undefined =
+          typeof value === 'string' ? undefined : value;
         attrSection.items.push({
           section: attrSection,
           label: key,
-          value: attrs[key],
+          value: strValue,
           canShowOnNode: true,
           showOnNode: this.curShowOnOpNodeAttrIds.has(key),
+          specialValue,
         });
       }
       if (attrSection.items.length > 0) {
@@ -1009,10 +1020,13 @@ export class InfoPanel {
     // Add tensor values to metadata if existed.
     const attrs = sourceOpNode.attrs || {};
     if (attrs[TENSOR_VALUES_KEY]) {
-      metadataList.push({
-        key: this.inputMetadataValuesKey,
-        value: attrs[TENSOR_VALUES_KEY],
-      });
+      const value = attrs[TENSOR_VALUES_KEY];
+      if (typeof value === 'string') {
+        metadataList.push({
+          key: this.inputMetadataValuesKey,
+          value,
+        });
+      }
     }
     return metadataList;
   }

src/ui/src/components/visualizer/io_tree.ts

@@ -252,7 +252,11 @@ export class IoTree implements OnChanges {
 
     if (isOpNode(modelNode)) {
       const attrs = modelNode.attrs || {};
-      return attrs[TENSOR_VALUES_KEY] || '<empty>';
+      const value = attrs[TENSOR_VALUES_KEY];
+      if (value && typeof value === 'string') {
+        return value;
+      }
+      return '<empty>';
     }
 
     return '';

src/ui/src/components/visualizer/worker/graph_processor.ts

@@ -31,6 +31,8 @@ import {
 import {
   KeyValuePairs,
   MetadataItem,
+  NodeAttributePairs,
+  NodeAttributeValue,
   NodeDataProviderRunData,
   ShowOnNodeItemData,
 } from '../common/types';
@@ -147,14 +149,18 @@ export class GraphProcessor {
           const attrValue = graphNode.attrs?.find(
             (attr) => attr.key === attrKey,
           )?.value;
-          if (attrValue && attrValue.match(attrValueRegex)) {
+          if (
+            attrValue &&
+            typeof attrValue === 'string' &&
+            attrValue.match(attrValueRegex)
+          ) {
             opNode.hideInLayout = true;
             break;
           }
         }
       }
       if (graphNode.attrs) {
-        const attrs: KeyValuePairs = {};
+        const attrs: NodeAttributePairs = {};
         for (const attr of graphNode.attrs) {
           attrs[attr.key] = this.processAttrValue(attr.key, attr.value);
         }
@@ -762,21 +768,28 @@ export class GraphProcessor {
     }
   }
 
-  private processAttrValue(key: string, value: string): string {
-    // Process const value that in `dense<...>` format. This is for backward
-    // compatibility.
-    if (value.startsWith('dense<')) {
-      const matches = value.match(CONST_VALUE_REGEX);
-      if (matches != null && matches.length > 1) {
-        const strTensorValue = matches[1];
-        return formatTensorValues(strTensorValue);
+  private processAttrValue(
+    key: string,
+    value: NodeAttributeValue,
+  ): NodeAttributeValue {
+    if (typeof value === 'string') {
+      // Process const value that in `dense<...>` format. This is for backward
+      // compatibility.
+      if (value.startsWith('dense<')) {
+        const matches = value.match(CONST_VALUE_REGEX);
+        if (matches != null && matches.length > 1) {
+          const strTensorValue = matches[1];
+          return formatTensorValues(strTensorValue);
+        }
       }
+      // Process tensor values.
+      else if (key === TENSOR_VALUES_KEY) {
+        return formatTensorValues(value);
+      }
+      return value.replaceAll('"', '') || '<empty>';
+    } else {
+      return value;
     }
-    // Process tensor values.
-    else if (key === TENSOR_VALUES_KEY) {
-      return formatTensorValues(value);
-    }
-    return value.replaceAll('"', '') || '<empty>';
   }
 
   private processMetadataList(metadataItems: MetadataItem[]) {