Switch edge scale to linear when the graph has no shape information.

Change: 119276088

tensorflow/tensorboard/components/tf-graph-common/lib/graph.ts

@@ -626,6 +626,7 @@ class MetaedgeImpl implements Metaedge {
     // Compute the size of the tensor flowing through this
     // base edge.
     this.totalSize += MetaedgeImpl.computeSizeOfEdge(edge, h);
+    h.maxMetaEdgeSize = Math.max(h.maxMetaEdgeSize, this.totalSize);
   }
 
   private static computeSizeOfEdge(edge: BaseEdge, h: hierarchy.Hierarchy):
@@ -636,6 +637,7 @@ class MetaedgeImpl implements Metaedge {
       // a lower bound for the total size.
       return 1;
     }
+    h.hasShapeInfo = true;
     // Sum the sizes of all output tensors.
     return _(opNode.outputShapes).map(shape => {
       // If the shape is unknown, treat it as 1 when computing
@@ -753,7 +755,7 @@ function extractOutputShapes(attr: {key: string, value: any}[]): TensorShape[] {
   }
   // We didn't find OUTPUT_SHAPES_KEY in attributes, so we don't know anything
   // about the output tensors.
-  return result;
+  return null;
 }
 
 /**

tensorflow/tensorboard/components/tf-graph-common/lib/hierarchy.ts

@@ -30,6 +30,10 @@ export interface Hierarchy {
   templates: {[templateId: string]: string[]};
   /** List of all device names */
   devices: string[];
+  /** True if at least one tensor in the graph has shape information */
+  hasShapeInfo: boolean;
+  /** The maximum size across all meta edges. Used for scaling thickness. */
+  maxMetaEdgeSize: number;
   getNodeMap(): {[nodeName: string]: GroupNode|OpNode};
   node(name: string): GroupNode|OpNode;
   setNode(name: string, node: GroupNode|OpNode): void;
@@ -48,6 +52,8 @@ class HierarchyImpl implements Hierarchy {
   templates: {[templateId: string]: string[]};
   private index: {[nodeName: string]: GroupNode|OpNode};
   devices: string[];
+  hasShapeInfo = false;
+  maxMetaEdgeSize = 1;
   orderings: { [nodeName: string]: { [childName: string]: number } };
 
   constructor() {

tensorflow/tensorboard/components/tf-graph-common/lib/render.ts

@@ -152,11 +152,14 @@ const PARAMS = {
  * for each node in the graph.
  */
 export class RenderGraphInfo {
-  private hierarchy: hierarchy.Hierarchy;
+  hierarchy: hierarchy.Hierarchy;
   private index: {[nodeName: string]: RenderNodeInfo};
   private deviceColorMap: d3.scale.Ordinal<string, string>;
   private memoryUsageScale: d3.scale.Linear<string, string>;
   private computeTimeScale: d3.scale.Linear<string, string>;
+  /** Scale for the thickness of edges when there is no shape information. */
+  edgeWidthScale:
+      d3.scale.Linear<number, number> | d3.scale.Pow<number, number>;
   // Since the rendering information for each node is constructed lazily,
   // upon node's expansion by the user, we keep a map between the node's name
   // and whether the rendering information was already constructed for that
@@ -210,6 +213,12 @@ export class RenderGraphInfo {
     this.computeTimeScale = d3.scale.linear<string, string>()
         .domain(computeTimeExtent)
         .range(PARAMS.minMaxColors);
+
+    this.edgeWidthScale = this.hierarchy.hasShapeInfo ?
+      scene.edge.EDGE_WIDTH_SCALE :
+      d3.scale.linear()
+        .domain([1, this.hierarchy.maxMetaEdgeSize])
+        .range([scene.edge.MIN_EDGE_WIDTH, scene.edge.MAX_EDGE_WIDTH]);
   }
 
   /**

tensorflow/tensorboard/components/tf-graph-common/lib/scene/edge.ts

@@ -18,10 +18,10 @@ module tf.graph.scene.edge {
 const TENSOR_SHAPE_DELIM = "×";
 
 /** The minimum stroke width of an edge. */
-const MIN_EDGE_WIDTH = 0.75;
+export const MIN_EDGE_WIDTH = 0.75;
 
 /** The maximum stroke width of an edge. */
-const MAX_EDGE_WIDTH = 12;
+export const MAX_EDGE_WIDTH = 12;
 
 /** The exponent used in the power scale for edge thickness. */
 const EDGE_WIDTH_SCALE_EXPONENT = 0.3;
@@ -29,7 +29,7 @@ const EDGE_WIDTH_SCALE_EXPONENT = 0.3;
 /** The domain (min and max value) for the edge width. */
 const DOMAIN_EDGE_WIDTH_SCALE = [1, 5E6];
 
-let edgeWidthScale = d3.scale.pow()
+export const EDGE_WIDTH_SCALE = d3.scale.pow()
       .exponent(EDGE_WIDTH_SCALE_EXPONENT)
       .domain(DOMAIN_EDGE_WIDTH_SCALE)
       .range([MIN_EDGE_WIDTH, MAX_EDGE_WIDTH])
@@ -229,7 +229,7 @@ export function appendEdge(edgeGroup, d: EdgeData,
   // Give the path a unique id, which will be used to link
   // the textPath (edge label) to this path.
   let pathId = "path_" + getEdgeKey(d);
-  let strokeWidth = edgeWidthScale(size);
+  let strokeWidth = sceneElement.renderHierarchy.edgeWidthScale(size);
 
   let path = edgeGroup.append("path")
     .attr({