Add a toggle that turns automatic high-degree node extraction off

tensorboard/plugins/graph/tf_graph/tf-graph.ts

@@ -135,6 +135,8 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
   renderHierarchy: tf_graph_render.RenderGraphInfo;
   @property({type: Boolean})
   traceInputs: boolean;
+  @property({type: Boolean})
+  autoExtractNodes: boolean;
   @property({type: Array})
   nodeContextMenuItems: unknown[];
   @property({
@@ -172,6 +174,15 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
   panToNode(nodeName) {
     (this.$$('tf-graph-scene') as any).panToNode(nodeName);
   }
+  @observe('autoExtractNodes')
+  _autoExtractNodesChanged() {
+    var graphHierarchy = this.graphHierarchy;
+    if (!graphHierarchy) return;
+    for (const node of Object.values(graphHierarchy.getNodeMap())) {
+      node.include = tf_graph.InclusionType.UNSPECIFIED;
+    }
+    this._buildRenderHierarchy(graphHierarchy);
+  }
   @observe(
     'graphHierarchy',
     'edgeWidthFunction',
@@ -264,9 +275,10 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
     const renderGraph = tf_graph_util.time(
       'new tf_graph_render.Hierarchy',
       () => {
-        var renderGraph = new tf_graph_render.RenderGraphInfo(
+        const renderGraph = new tf_graph_render.RenderGraphInfo(
           graphHierarchy,
-          !!this.stats /** displayingStats */
+          !!this.stats /** displayingStats */,
+          this.autoExtractNodes
         );
         renderGraph.edgeLabelFunction = this.edgeLabelFunction;
         renderGraph.edgeWidthFunction = this.edgeWidthFunction;

tensorboard/plugins/graph/tf_graph_app/tf-graph-app.ts

@@ -97,6 +97,7 @@ class TfGraphApp extends LegacyElementMixin(PolymerElement) {
             selected-file="{{selectedFile}}"
             on-fit-tap="_fit"
             trace-inputs="{{_traceInputs}}"
+            auto-extract-nodes="{{_autoExtractNodes}}"
           ></tf-graph-controls>
           <tf-graph-loader
             id="loader"
@@ -117,6 +118,7 @@ class TfGraphApp extends LegacyElementMixin(PolymerElement) {
             render-hierarchy="{{_renderHierarchy}}"
             selected-node="{{selectedNode}}"
             trace-inputs="[[_traceInputs]]"
+            autoExtractNodes="[[_autoExtractNodes]]"
           ></tf-graph-board>
         </div>
       </div>
@@ -164,6 +166,8 @@ class TfGraphApp extends LegacyElementMixin(PolymerElement) {
   _progress: object;
   @property({type: Boolean})
   _traceInputs: boolean;
+  @property({type: Boolean})
+  _autoExtractNodes: boolean;
   _updateToolbar() {
     (this.$$('.container') as HTMLElement).classList.toggle(
       'no-toolbar',

tensorboard/plugins/graph/tf_graph_board/tf-graph-board.ts

@@ -152,6 +152,7 @@ class TfGraphBoard extends LegacyElementMixin(PolymerElement) {
           handle-node-selected="[[handleNodeSelected]]"
           handle-edge-selected="[[handleEdgeSelected]]"
           trace-inputs="[[traceInputs]]"
+          auto-extract-nodes="[[autoExtractNodes]]"
         ></tf-graph>
       </div>
       <div id="info">
@@ -196,6 +197,8 @@ class TfGraphBoard extends LegacyElementMixin(PolymerElement) {
   progress: object;
   @property({type: Boolean})
   traceInputs: boolean;
+  @property({type: Boolean})
+  autoExtractNodes: boolean;
   @property({type: String})
   colorBy: string;
   @property({

tensorboard/plugins/graph/tf_graph_common/render.ts

@@ -127,10 +127,6 @@ export interface EdgeLabelFunction {
  * Parameters that affect how the graph is rendered on the screen.
  */
 const PARAMS = {
-  /**
-   * Whether to extract high degree nodes from the core part of the graph.
-   */
-  enableExtraction: true,
   /**
    * The minimum number of nodes for a graph to have in order for high in and
    * out degree nodes to be extracted in auxiliary. The aim here is to prevent
@@ -233,14 +229,20 @@ export class RenderGraphInfo {
     [nodeName: string]: boolean;
   };
   root: RenderGroupNodeInfo;
-  traceInputs: Boolean;
+  traceInputs: boolean;
+  private autoExtractNodes: boolean;
   edgeLabelFunction: EdgeLabelFunction;
   // An optional function that computes the thickness of an edge given edge
   // data. If not provided, defaults to encoding tensor size in thickness.
   edgeWidthFunction: EdgeThicknessFunction;
-  constructor(hierarchy: tf_graph.Hierarchy, displayingStats: boolean) {
+  constructor(
+    hierarchy: tf_graph.Hierarchy,
+    displayingStats: boolean,
+    autoExtractNodes: boolean
+  ) {
     this.hierarchy = hierarchy;
     this.displayingStats = displayingStats;
+    this.autoExtractNodes = autoExtractNodes;
     this.index = {};
     this.renderedOpNames = [];
     this.computeScales();
@@ -942,11 +944,8 @@ export class RenderGraphInfo {
         this.index[edgeObj.v].isFadedOut || this.index[edgeObj.w].isFadedOut;
       coreGraph.setEdge(edgeObj.v, edgeObj.w, renderMetaedgeInfo);
     });
-    if (
-      PARAMS.enableExtraction &&
-      renderGroupNodeInfo.node.type === NodeType.META
-    ) {
-      extractHighDegrees(renderGroupNodeInfo);
+    if (renderGroupNodeInfo.node.type === NodeType.META) {
+      extractHighDegrees(renderGroupNodeInfo, this.autoExtractNodes);
     }
     // If there are functions, it is possible for metanodes to be dynamically
     // added later. Construct the hierarchies for nodes that are predecessors to
@@ -2182,18 +2181,29 @@ export function mapIndexToHue(id: number): number {
  * screw up the graph layout.
  *
  * @param {Render.Node} renderNode Node to manipulate.
+ * @param {boolean} autoExtractNodes Whether to automatically exclude high-degree
+ *   nodes from the main graph. If false, only exclude predefined nodes.
  */
-function extractHighDegrees(renderNode: RenderGroupNodeInfo) {
+function extractHighDegrees(
+  renderNode: RenderGroupNodeInfo,
+  autoExtractNodes: boolean
+) {
   extractSpecifiedNodes(renderNode);
+
   if (PARAMS.outExtractTypes.length) {
     extractPredefinedSink(renderNode);
   }
+
   // This has to come before extract high in-degree to protect the core part
   // that takes many variables.
   if (PARAMS.inExtractTypes.length) {
     extractPredefinedSource(renderNode);
   }
-  extractHighInOrOutDegree(renderNode);
+
+  if (autoExtractNodes) {
+    extractHighInOrOutDegree(renderNode);
+  }
+
   if (PARAMS.maxControlDegree) {
     removeControlEdges(renderNode);
   }

tensorboard/plugins/graph/tf_graph_controls/tf-graph-controls.ts

@@ -492,6 +492,13 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
           </paper-toggle-button>
         </div>
       </div>
+      <div class="control-holder">
+        <div>
+          <paper-toggle-button checked="{{autoExtractNodes}}" class="title">
+            Auto-extract high-degree nodes
+          </paper-toggle-button>
+        </div>
+      </div>
       <template is="dom-if" if="[[healthPillsFeatureEnabled]]">
         <div class="control-holder">
           <paper-toggle-button checked="{{healthPillsToggledOn}}" class="title"
@@ -1079,6 +1086,11 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
     notify: true,
   })
   traceInputs: boolean = false;
+  @property({
+    type: Boolean,
+    notify: true,
+  })
+  autoExtractNodes: boolean = true;
   @property({
     type: Number,
     observer: '_selectedTagIndexChanged',

tensorboard/plugins/graph/tf_graph_dashboard/tf-graph-dashboard.ts

@@ -87,6 +87,7 @@ class TfGraphDashboard extends LegacyElementMixin(PolymerElement) {
         health-pills-toggled-on="{{healthPillsToggledOn}}"
         on-fit-tap="_fit"
         trace-inputs="{{_traceInputs}}"
+        auto-extract-nodes="{{_autoExtractNodes}}"
       ></tf-graph-controls>
       <div
         class$="center [[_getGraphDisplayClassName(_selectedFile, _datasets)]]"
@@ -160,6 +161,7 @@ class TfGraphDashboard extends LegacyElementMixin(PolymerElement) {
             selected-node="{{_selectedNode}}"
             stats="[[_stats]]"
             trace-inputs="[[_traceInputs]]"
+            auto-extract-nodes="[[_autoExtractNodes]]"
           ></tf-graph-board>
         </div>
       </div>
@@ -304,6 +306,8 @@ class TfGraphDashboard extends LegacyElementMixin(PolymerElement) {
   _compatibilityProvider: object;
   @property({type: Boolean})
   _traceInputs: boolean;
+  @property({type: Boolean})
+  _autoExtractNodes: boolean;
   @property({type: Object})
   _selectedFile: any;
   attached() {