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

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

@@ -44,6 +44,8 @@ class TfGraphScene2 extends LegacyElementMixin(PolymerElement) {
   colorBy: string;
   @property({type: Boolean})
   traceInputs: boolean;
+  @property({type: Boolean})
+  extractNodes: boolean;
 
   // For each render hierarchy, we only fit it to the viewport once (when the scene is attached to
   // the DOM). We do not fit the hierarchy again (unless the user clicks the reset button). For
@@ -663,4 +665,10 @@ class TfGraphScene2 extends LegacyElementMixin(PolymerElement) {
       this.traceInputs
     );
   }
+
+  // Re-do hierarchy if extraction changes
+  @observe('extractNodes')
+  _updateExtractNodes() {
+    this.renderHierarchy
+  }
 }

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

@@ -77,6 +77,7 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
           health-pill-step-index="{{healthPillStepIndex}}"
           handle-edge-selected="[[handleEdgeSelected]]"
           trace-inputs="[[traceInputs]]"
+          extract-nodes="[[extractNodes]]"
         ></tf-graph-scene>
       </div>
     </div>
@@ -135,6 +136,8 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
   renderHierarchy: tf_graph_render.RenderGraphInfo;
   @property({type: Boolean})
   traceInputs: boolean;
+  @property({type: Boolean})
+  extractNodes: boolean;
   @property({type: Array})
   nodeContextMenuItems: unknown[];
   @property({
@@ -173,6 +176,7 @@ class TfGraph extends LegacyElementMixin(PolymerElement) {
     (this.$$('tf-graph-scene') as any).panToNode(nodeName);
   }
   @observe(
+    'extractNodes',
     'graphHierarchy',
     'edgeWidthFunction',
     'handleNodeSelected',
@@ -264,9 +268,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.extractNodes
         );
         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}}"
+            extract-nodes="{{_extractNodes}}"
           ></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]]"
+            extractNodes="[[_extractNodes]]"
           ></tf-graph-board>
         </div>
       </div>
@@ -164,6 +166,8 @@ class TfGraphApp extends LegacyElementMixin(PolymerElement) {
   _progress: object;
   @property({type: Boolean})
   _traceInputs: boolean;
+  @property({type: Boolean})
+  _extractNodes: 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]]"
+          extract-nodes="[[extractNodes]]"
         ></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})
+  extractNodes: 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,18 @@ export class RenderGraphInfo {
     [nodeName: string]: boolean;
   };
   root: RenderGroupNodeInfo;
-  traceInputs: Boolean;
+  traceInputs: boolean;
+  extractNodes: 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,
+              extractNodes: boolean) {
     this.hierarchy = hierarchy;
     this.displayingStats = displayingStats;
+    this.extractNodes = extractNodes;
     this.index = {};
     this.renderedOpNames = [];
     this.computeScales();
@@ -942,11 +942,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.extractNodes);
     }
     // 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 +2179,27 @@ export function mapIndexToHue(id: number): number {
  * screw up the graph layout.
  *
  * @param {Render.Node} renderNode Node to manipulate.
+ * @param {boolean} extractNodes 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,
+                            extractNodes: 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 (extractNodes) {
+    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="{{extractNodes}}" 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,
+  })
+  extractNodes: boolean = false;
   @property({
     type: Number,
     observer: '_selectedTagIndexChanged',
@@ -1357,6 +1369,7 @@ class TfGraphControls extends LegacyElementMixin(PolymerElement) {
     this._selectedTagIndex = 0;
     this._selectedGraphType = this._getDefaultSelectionType();
     this.traceInputs = false; // Set trace input to off-state.
+    this.extractNodes = true;
     this._setDownloadFilename(
       this.datasets[runIndex] ? this.datasets[runIndex].name : ''
     );

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}}"
+        extract-nodes="{{_extractNodes}}"
       ></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]]"
+            extract-nodes="[[_extractNodes]]"
           ></tf-graph-board>
         </div>
       </div>
@@ -304,6 +306,8 @@ class TfGraphDashboard extends LegacyElementMixin(PolymerElement) {
   _compatibilityProvider: object;
   @property({type: Boolean})
   _traceInputs: boolean;
+  @property({type: Boolean})
+  _extractNodes: boolean;
   @property({type: Object})
   _selectedFile: any;
   attached() {