Adding tensor shapes to the graph visualizer and to the info card.

Also:
- rename sceneBehavior -> sceneElement to make it clearer it is a polymer element.
- improve the info card by showing the actual op node in the successors/predecessors list when the metaedge only contains one base edge (one tensor).
Change: 115339805

5 files changed, 200 insertions, 127 deletions

diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/annotation.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/annotation.ts
index 425eea0408..4301acc8d9 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/annotation.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/annotation.ts
@@ -39,11 +39,11 @@ module tf.graph.scene.annotation {
  * @param container selection of the container.
  * @param annotationData node.{in|out}Annotations
  * @param d node to build group for.
- * @param sceneBehavior polymer scene element.
+ * @param sceneElement <tf-graph-scene> polymer element.
  * @return selection of appended objects
  */
 export function buildGroup(container, annotationData: render.AnnotationList,
-  d: render.RenderNodeInfo, sceneBehavior) {
+  d: render.RenderNodeInfo, sceneElement) {
   // Select all children and join with data.
   let annotationGroups = container.selectAll(function() {
        // using d3's selector function
@@ -60,7 +60,7 @@ export function buildGroup(container, annotationData: render.AnnotationList,
       let aGroup = d3.select(this);
 
       // Add annotation to the index in the scene
-      sceneBehavior.addAnnotationGroup(a, d, aGroup);
+      sceneElement.addAnnotationGroup(a, d, aGroup);
       // Append annotation edge
       let edgeType = Class.Annotation.EDGE;
       let metaedge = a.renderMetaedgeInfo && a.renderMetaedgeInfo.metaedge;
@@ -71,11 +71,11 @@ export function buildGroup(container, annotationData: render.AnnotationList,
       if (metaedge && metaedge.numRefEdges) {
         edgeType += " " + Class.Edge.REF_LINE;
       }
-      edge.appendEdge(aGroup, a, sceneBehavior, edgeType);
+      edge.appendEdge(aGroup, a, sceneElement, edgeType);
 
-      if (a.annotationType !== tf.graph.render.AnnotationType.ELLIPSIS) {
+      if (a.annotationType !== render.AnnotationType.ELLIPSIS) {
         addAnnotationLabelFromNode(aGroup, a);
-        buildShape(aGroup, a, sceneBehavior);
+        buildShape(aGroup, a);
       } else {
         addAnnotationLabel(aGroup, a.node.name, a, Class.Annotation.ELLIPSIS);
       }
@@ -89,9 +89,9 @@ export function buildGroup(container, annotationData: render.AnnotationList,
     })
     .each(function(a) {
       let aGroup = d3.select(this);
-      update(aGroup, d, a, sceneBehavior);
-      if (a.annotationType !== tf.graph.render.AnnotationType.ELLIPSIS) {
-        addInteraction(aGroup, d, a, sceneBehavior);
+      update(aGroup, d, a, sceneElement);
+      if (a.annotationType !== render.AnnotationType.ELLIPSIS) {
+        addInteraction(aGroup, d, a, sceneElement);
       }
     });
 
@@ -100,7 +100,7 @@ export function buildGroup(container, annotationData: render.AnnotationList,
       let aGroup = d3.select(this);
 
       // Remove annotation from the index in the scene
-      sceneBehavior.removeAnnotationGroup(a, d, aGroup);
+      sceneElement.removeAnnotationGroup(a, d, aGroup);
     })
     .remove();
   return annotationGroups;
@@ -110,13 +110,13 @@ export function buildGroup(container, annotationData: render.AnnotationList,
  * Maps an annotation enum to a class name used in css rules.
  */
 function annotationToClassName(annotationType: render.AnnotationType) {
-  return (tf.graph.render.AnnotationType[annotationType] || "")
+  return (render.AnnotationType[annotationType] || "")
       .toLowerCase() || null;
 }
 
-function buildShape(aGroup, a: render.Annotation, sceneBehavior) {
-  if (a.annotationType === tf.graph.render.AnnotationType.SUMMARY) {
-    let summary = scene.selectOrCreateChild(aGroup, "use");
+function buildShape(aGroup, a: render.Annotation) {
+  if (a.annotationType === render.AnnotationType.SUMMARY) {
+    let summary = selectOrCreateChild(aGroup, "use");
     summary.attr({
       "class": "summary",
       "xlink:href": "#summary-icon",
@@ -125,7 +125,7 @@ function buildShape(aGroup, a: render.Annotation, sceneBehavior) {
   } else {
     let shape = node.buildShape(aGroup, a, Class.Annotation.NODE);
     // add title tag to get native tooltips
-    scene.selectOrCreateChild(shape, "title").text(a.node.name);
+    selectOrCreateChild(shape, "title").text(a.node.name);
   }
 }
 
@@ -152,16 +152,16 @@ function addAnnotationLabel(aGroup, label, a, additionalClassNames,
 }
 
 function addInteraction(selection, d: render.RenderNodeInfo,
-    annotation: tf.graph.render.Annotation, sceneBehavior) {
+    annotation: render.Annotation, sceneElement) {
   selection
     .on("mouseover", a => {
-      sceneBehavior.fire("annotation-highlight", {
+      sceneElement.fire("annotation-highlight", {
         name: a.node.name,
         hostName: d.node.name
       });
     })
     .on("mouseout", a => {
-      sceneBehavior.fire("annotation-unhighlight", {
+      sceneElement.fire("annotation-unhighlight", {
         name: a.node.name,
         hostName: d.node.name
       });
@@ -170,15 +170,15 @@ function addInteraction(selection, d: render.RenderNodeInfo,
       // Stop this event"s propagation so that it isn't also considered a
       // graph-select.
       (<Event>d3.event).stopPropagation();
-      sceneBehavior.fire("annotation-select", {
+      sceneElement.fire("annotation-select", {
         name: a.node.name,
         hostName: d.node.name
       });
     });
-  if (annotation.annotationType !== tf.graph.render.AnnotationType.SUMMARY &&
-      annotation.annotationType !== tf.graph.render.AnnotationType.CONSTANT) {
-    selection.on("contextmenu", tf.graph.scene.contextmenu.getMenu(
-      tf.graph.scene.node.getContextMenu(annotation.node, sceneBehavior)));
+  if (annotation.annotationType !== render.AnnotationType.SUMMARY &&
+      annotation.annotationType !== render.AnnotationType.CONSTANT) {
+    selection.on("contextmenu", contextmenu.getMenu(
+      node.getContextMenu(annotation.node, sceneElement)));
   }
 };
 
@@ -188,21 +188,21 @@ function addInteraction(selection, d: render.RenderNodeInfo,
  * @param aGroup selection of a "g.annotation" element.
  * @param d Host node data.
  * @param a annotation node data.
- * @param scene Polymer scene element.
+ * @param scene <tf-graph-scene> polymer element.
  */
 function update(aGroup, d: render.RenderNodeInfo, a: render.Annotation,
-    sceneBehavior) {
+    sceneElement) {
   let cx = layout.computeCXPositionOfNodeShape(d);
   // Annotations that point to embedded nodes (constants,summary)
   // don't have a render information attached so we don't stylize these.
   // Also we don't stylize ellipsis annotations (the string "... and X more").
   if (a.renderNodeInfo &&
-      a.annotationType !== tf.graph.render.AnnotationType.ELLIPSIS) {
-    node.stylize(aGroup, a.renderNodeInfo, sceneBehavior,
+      a.annotationType !== render.AnnotationType.ELLIPSIS) {
+    node.stylize(aGroup, a.renderNodeInfo, sceneElement,
       Class.Annotation.NODE);
   }
 
-  if (a.annotationType === tf.graph.render.AnnotationType.SUMMARY) {
+  if (a.annotationType === render.AnnotationType.SUMMARY) {
     // Update the width of the annotation to give space for the image.
     a.width += 10;
   }
@@ -224,11 +224,11 @@ function update(aGroup, d: render.RenderNodeInfo, a: render.Annotation,
   });
 
   // Node position (only one of the shape selection will be non-empty.)
-  scene.positionEllipse(aGroup.select("." + Class.Annotation.NODE + " ellipse"),
+  positionEllipse(aGroup.select("." + Class.Annotation.NODE + " ellipse"),
                         cx + a.dx, d.y + a.dy, a.width, a.height);
-  scene.positionRect(aGroup.select("." + Class.Annotation.NODE + " rect"),
+  positionRect(aGroup.select("." + Class.Annotation.NODE + " rect"),
                      cx + a.dx, d.y + a.dy, a.width, a.height);
-  scene.positionRect(aGroup.select("." + Class.Annotation.NODE + " use"),
+  positionRect(aGroup.select("." + Class.Annotation.NODE + " use"),
                      cx + a.dx, d.y + a.dy, a.width, a.height);
 
   // Edge position
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/edge.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/edge.ts
index 5ae6244e00..d0f1e8fad6 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/edge.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/edge.ts
@@ -19,10 +19,13 @@ limitations under the License.
 
 module tf.graph.scene.edge {
 
+/** Delimiter between dimensions when showing sizes of tensors. */
+const TENSOR_SHAPE_DELIM = "×";
+
 export type EdgeData = {v: string, w: string, label: render.RenderMetaedgeInfo};
 
 export function getEdgeKey(edgeObj: EdgeData) {
-  return edgeObj.v + tf.graph.EDGE_KEY_DELIM + edgeObj.w;
+  return edgeObj.v + EDGE_KEY_DELIM + edgeObj.w;
 }
 
 /**
@@ -41,14 +44,14 @@ export function getEdgeKey(edgeObj: EdgeData) {
  *
  * @param sceneGroup container
  * @param graph
- * @param sceneBehavior Parent scene module.
+ * @param sceneElement <tf-graph-scene> polymer element.
  * @return selection of the created nodeGroups
  */
 export function buildGroup(sceneGroup,
-  graph: graphlib.Graph<tf.graph.render.RenderNodeInfo,
-    tf.graph.render.RenderMetaedgeInfo>, sceneBehavior) {
+    graph: graphlib.Graph<render.RenderNodeInfo, render.RenderMetaedgeInfo>,
+    sceneElement) {
   let edges: EdgeData[] = [];
-  let edgeData = _.reduce(graph.edges(), (edges, edgeObj) => {
+  edges = _.reduce(graph.edges(), (edges, edgeObj) => {
     let edgeLabel = graph.edge(edgeObj);
     edges.push({
       v: edgeObj.v,
@@ -68,8 +71,7 @@ export function buildGroup(sceneGroup,
     // See https://github.com/mbostock/d3/releases/tag/v2.0.0
     // (It's not listed in the d3 wiki.)
     return this.childNodes;
-  })
-    .data(edgeData, getEdgeKey);
+  }).data(edges, getEdgeKey);
 
   // Make edges a group to support rendering multiple lines for metaedge
   edgeGroups.enter()
@@ -80,7 +82,7 @@ export function buildGroup(sceneGroup,
       let edgeGroup = d3.select(this);
       d.label.edgeGroup = edgeGroup;
       // index node group for quick highlighting
-      sceneBehavior._edgeGroupIndex[getEdgeKey(d)] = edgeGroup;
+      sceneElement._edgeGroupIndex[getEdgeKey(d)] = edgeGroup;
 
       // If any edges are reference edges, add the reference edge class.
       let extraEdgeClass = d.label.metaedge && d.label.metaedge.numRefEdges
@@ -88,22 +90,57 @@ export function buildGroup(sceneGroup,
         : undefined;
       // Add line during enter because we're assuming that type of line
       // normally does not change.
-      appendEdge(edgeGroup, d, scene, extraEdgeClass);
+      appendEdge(edgeGroup, d, sceneElement, extraEdgeClass);
     });
 
   edgeGroups.each(position);
   edgeGroups.each(function(d) {
-    stylize(d3.select(this), d, sceneBehavior);
+    stylize(d3.select(this), d, sceneElement);
   });
 
   edgeGroups.exit()
     .each(d => {
-      delete sceneBehavior._edgeGroupIndex[getEdgeKey(d)];
+      delete sceneElement._edgeGroupIndex[getEdgeKey(d)];
     })
     .remove();
   return edgeGroups;
 };
 
+export function getShapeLabelFromNode(node: OpNode,
+    renderInfo: render.RenderGraphInfo) {
+  if (node.outputShapes == null || node.outputShapes.length === 0) {
+    return null;
+  }
+  // TODO(smilkov): Figure out exactly which output tensor this
+  // edge is from.
+  let shape = node.outputShapes[0];
+  if (shape == null) {
+    return null;
+  }
+  if (shape.length === 0) {
+    return "scalar";
+  }
+  return shape.map(size => {
+    return size === -1 ? "?" : size;
+  }).join(TENSOR_SHAPE_DELIM);
+}
+
+/**
+ * Creates the label for the given metaedge. If the metaedge consists
+ * of only 1 tensor, and it's shape is known, the label will contain that
+ * shape. Otherwise, the label will say the number of tensors in the metaedge.
+ */
+export function getLabelForEdge(metaedge: Metaedge,
+    renderInfo: render.RenderGraphInfo): string {
+  let isMultiEdge = metaedge.baseEdgeList.length > 1;
+  if (isMultiEdge) {
+    return metaedge.baseEdgeList.length + " tensors";
+  } else {
+    let node = <OpNode> renderInfo.getNodeByName(metaedge.baseEdgeList[0].v);
+    return getShapeLabelFromNode(node, renderInfo);
+  }
+}
+
 /**
  * For a given d3 selection and data object, create a path to represent the
  * edge described in d.label.
@@ -112,15 +149,49 @@ export function buildGroup(sceneGroup,
  * will sometimes be undefined, for example for some Annotation edges for which
  * there is no underlying Metaedge in the hierarchical graph.
  */
-export function appendEdge(edgeGroup, d: EdgeData, sceneBehavior, edgeClass?) {
+export function appendEdge(edgeGroup, d: EdgeData,
+    sceneElement: {renderHierarchy: render.RenderGraphInfo},
+    edgeClass: string) {
+  let size = 1;
+  if (d.label != null && d.label.metaedge != null) {
+    // There is an underlying Metaedge.
+    size = d.label.metaedge.totalSize;
+  }
   edgeClass = edgeClass || Class.Edge.LINE; // set default type
 
   if (d.label && d.label.structural) {
     edgeClass += " " + Class.Edge.STRUCTURAL;
   }
+  // 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 = sceneElement.renderHierarchy.edgeWidthScale(size);
 
   edgeGroup.append("path")
-    .attr("class", edgeClass);
+    .attr({
+      "id": pathId,
+      "class": edgeClass,
+    }).style({
+      "stroke-width": strokeWidth + "px"
+    });
+
+  if (d.label == null || d.label.metaedge == null) {
+    // There is no associated metaedge, thus no text.
+    // This happens for annotation edges.
+    return;
+  }
+  let labelForEdge = getLabelForEdge(d.label.metaedge,
+      sceneElement.renderHierarchy);
+  if (labelForEdge == null) {
+    // We have no information to show on this edge.
+    return;
+  }
+  edgeGroup.append("text").append("textPath").attr({
+      "xlink:href": "#" + pathId,
+      "startOffset": "50%",
+      "text-anchor": "middle",
+      "dominant-baseline": "central"
+  }).text(labelForEdge);
 };
 
 export let interpolate = d3.svg.line<{x: number, y: number}>()
@@ -134,8 +205,9 @@ export let interpolate = d3.svg.line<{x: number, y: number}>()
 function getEdgePathInterpolator(d: EdgeData, i: number, a: string) {
   let renderMetaedgeInfo = <render.RenderMetaedgeInfo> d.label;
   let adjoiningMetaedge = renderMetaedgeInfo.adjoiningMetaedge;
+  let points = renderMetaedgeInfo.points;
   if (!adjoiningMetaedge) {
-    return d3.interpolate(a, interpolate(renderMetaedgeInfo.points));
+    return d3.interpolate(a, interpolate(points));
   }
 
   let renderPath = this;
@@ -158,7 +230,6 @@ function getEdgePathInterpolator(d: EdgeData, i: number, a: string) {
 
     // Update the relevant point in the renderMetaedgeInfo's points list, then
     // re-interpolate the path.
-    let points = renderMetaedgeInfo.points;
     let index = inbound ? 0 : points.length - 1;
     points[index].x = adjoiningPoint.x;
     points[index].y = adjoiningPoint.y;
@@ -169,10 +240,8 @@ function getEdgePathInterpolator(d: EdgeData, i: number, a: string) {
 
 function position(d) {
   d3.select(this).select("path." + Class.Edge.LINE)
-    .each(function(d) {
-      let path = d3.select(this);
-      path.transition().attrTween("d", getEdgePathInterpolator);
-    });
+    .transition()
+    .attrTween("d", getEdgePathInterpolator);
 };
 
 /**
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/minimap.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/minimap.ts
index 32566c99ef..0be15da0e0 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/minimap.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/minimap.ts
@@ -262,7 +262,8 @@ export class Minimap {
       downloadContext.drawImage(image, 0, 0,
         this.downloadCanvas.width, this.downloadCanvas.height);
     };
-    image.src = "data:image/svg+xml;base64," + btoa(svgXml);
+    let blob = new Blob([svgXml], {type: "image/svg+xml;charset=utf-8"});
+    image.src = URL.createObjectURL(blob);
   }
 
   /**
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/node.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/node.ts
index 58ec821877..6d467c348a 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/node.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/node.ts
@@ -61,11 +61,11 @@ module tf.graph.scene.node {
  *
  * @param sceneGroup selection of the container
  * @param nodeData array of render node information to map
- * @param sceneBehavior parent scene module
+ * @param sceneElement <tf-graph-scene> polymer element
  * @return selection of the created nodeGroups
  */
 export function buildGroup(sceneGroup,
-    nodeData: render.RenderNodeInfo[], sceneBehavior) {
+    nodeData: render.RenderNodeInfo[], sceneElement) {
   let container = scene.selectOrCreateChild(sceneGroup, "g",
     Class.Node.CONTAINER);
   // Select all children and join with data.
@@ -88,7 +88,7 @@ export function buildGroup(sceneGroup,
     .each(function(d) {
       let nodeGroup = d3.select(this);
       // index node group for quick stylizing
-      sceneBehavior.addNodeGroup(d.node.name, nodeGroup);
+      sceneElement.addNodeGroup(d.node.name, nodeGroup);
     });
 
   // UPDATE
@@ -98,58 +98,58 @@ export function buildGroup(sceneGroup,
     })
     .each(function(d) {
       let nodeGroup = d3.select(this);
-      // add g.in-annotations (always add -- to keep layer order consistent.)
+      // Add g.in-annotations (always add -- to keep layer order consistent.)
       let inAnnotationBox = scene.selectOrCreateChild(nodeGroup, "g",
         Class.Annotation.INBOX);
       annotation.buildGroup(inAnnotationBox, d.inAnnotations, d,
-        sceneBehavior);
+        sceneElement);
 
-      // add g.out-annotations  (always add -- to keep layer order consistent.)
+      // Add g.out-annotations  (always add -- to keep layer order consistent.)
       let outAnnotationBox = scene.selectOrCreateChild(nodeGroup, "g",
         Class.Annotation.OUTBOX);
       annotation.buildGroup(outAnnotationBox, d.outAnnotations, d,
-        sceneBehavior);
+        sceneElement);
 
-      // label
-      let label = labelBuild(nodeGroup, d, sceneBehavior);
-      // Do not add interaction to metanode labels as they live inside the
-      // metanode shape which already has the same interactions.
-      addInteraction(label, d, sceneBehavior, d.node.type === NodeType.META);
-
-      // build .shape below label
-      let shape = buildShape(nodeGroup, d, Class.Node.SHAPE, label.node());
+      // Build .shape first (background of the node).
+      let shape = buildShape(nodeGroup, d, Class.Node.SHAPE);
       if (d.node.isGroupNode) {
-        addButton(shape, d, sceneBehavior);
+        addButton(shape, d, sceneElement);
       }
-      addInteraction(shape, d, sceneBehavior);
+      addInteraction(shape, d, sceneElement);
 
-      // build subscene on the top
+      // Build subscene on the top.
       subsceneBuild(nodeGroup, <render.RenderGroupNodeInfo> d,
-          sceneBehavior);
+          sceneElement);
+
+      // Build label last. Should be on top of everything else.
+      let label = labelBuild(nodeGroup, d, sceneElement);
+      // Do not add interaction to metanode labels as they live inside the
+      // metanode shape which already has the same interactions.
+      addInteraction(label, d, sceneElement, d.node.type === NodeType.META);
 
-      stylize(nodeGroup, d, sceneBehavior);
-      position(nodeGroup, d, sceneBehavior);
+      stylize(nodeGroup, d, sceneElement);
+      position(nodeGroup, d);
     });
 
   // EXIT
   nodeGroups.exit()
     .each(function(d) {
       // remove all indices on remove
-      sceneBehavior.removeNodeGroup(d.node.name);
+      sceneElement.removeNodeGroup(d.node.name);
 
       let nodeGroup = d3.select(this);
       if (d.inAnnotations.list.length > 0) {
         nodeGroup.select("." + Class.Annotation.INBOX)
           .selectAll("." + Class.Annotation.GROUP)
           .each(a => {
-            sceneBehavior.removeAnnotationGroup(a, d);
+            sceneElement.removeAnnotationGroup(a, d);
           });
       }
       if (d.outAnnotations.list.length > 0) {
         nodeGroup.select("." + Class.Annotation.OUTBOX)
           .selectAll("." + Class.Annotation.GROUP)
           .each(a => {
-            sceneBehavior.removeAnnotationGroup(a, d);
+            sceneElement.removeAnnotationGroup(a, d);
           });
       }
     })
@@ -163,17 +163,17 @@ export function buildGroup(sceneGroup,
  *
  * @param nodeGroup selection of the container
  * @param renderNodeInfo the render information for the node.
- * @param sceneBehavior parent scene module
+ * @param sceneElement <tf-graph-scene> polymer element.
  * @return Selection of the subscene group, or null if node group does not have
  *        a subscene. Op nodes, bridge nodes and unexpanded group nodes will
  *        not have a subscene.
  */
 function subsceneBuild(nodeGroup,
-    renderNodeInfo: render.RenderGroupNodeInfo, sceneBehavior) {
+    renderNodeInfo: render.RenderGroupNodeInfo, sceneElement) {
   if (renderNodeInfo.node.isGroupNode) {
     if (renderNodeInfo.expanded) {
       // Recursively build the subscene.
-      return scene.buildGroup(nodeGroup, renderNodeInfo, sceneBehavior,
+      return scene.buildGroup(nodeGroup, renderNodeInfo, sceneElement,
         Class.Subscene.GROUP);
     }
     // Clean out existing subscene if the node is not expanded.
@@ -198,9 +198,9 @@ function subscenePosition(nodeGroup, d: render.RenderNodeInfo) {
  *
  * @param selection The group node selection.
  * @param d Info about the node being rendered.
- * @param sceneBehavior parent scene module.
+ * @param sceneElement <tf-graph-scene> polymer element.
  */
-function addButton(selection, d: render.RenderNodeInfo, sceneBehavior) {
+function addButton(selection, d: render.RenderNodeInfo, sceneElement) {
   let group = scene.selectOrCreateChild(
     selection, "g", Class.Node.BUTTON_CONTAINER);
   scene.selectOrCreateChild(group, "circle", Class.Node.BUTTON_CIRCLE);
@@ -212,7 +212,7 @@ function addButton(selection, d: render.RenderNodeInfo, sceneBehavior) {
     // Stop this event's propagation so that it isn't also considered a
     // node-select.
     (<Event>d3.event).stopPropagation();
-    sceneBehavior.fire("node-toggle-expand", { name: d.node.name });
+    sceneElement.fire("node-toggle-expand", { name: d.node.name });
   });
   scene.positionButton(group, d);
 };
@@ -226,39 +226,39 @@ function addButton(selection, d: render.RenderNodeInfo, sceneBehavior) {
  * given interaction would cause conflicts with the expand/collapse button.
  */
 function addInteraction(selection, d: render.RenderNodeInfo,
-    sceneBehavior, disableInteraction?: boolean) {
+    sceneElement, disableInteraction?: boolean) {
   if (disableInteraction) {
     selection.attr("pointer-events", "none");
     return;
   }
 
-  let contextMenuFunction = tf.graph.scene.contextmenu.getMenu(
-    getContextMenu(d.node, sceneBehavior));
+  let contextMenuFunction = contextmenu.getMenu(
+    getContextMenu(d.node, sceneElement));
   selection.on("dblclick", d => {
-    sceneBehavior.fire("node-toggle-expand", { name: d.node.name });
+    sceneElement.fire("node-toggle-expand", { name: d.node.name });
   })
     .on("mouseover", d => {
       // don't send mouseover over expanded group,
       // otherwise it is causing too much glitches
-      if (sceneBehavior.isNodeExpanded(d)) { return; }
+      if (sceneElement.isNodeExpanded(d)) { return; }
 
-      sceneBehavior.fire("node-highlight", { name: d.node.name });
+      sceneElement.fire("node-highlight", { name: d.node.name });
     })
     .on("mouseout", d => {
       // don't send mouseover over expanded group,
       // otherwise it is causing too much glitches
-      if (sceneBehavior.isNodeExpanded(d)) { return; }
+      if (sceneElement.isNodeExpanded(d)) { return; }
 
-      sceneBehavior.fire("node-unhighlight", { name: d.node.name });
+      sceneElement.fire("node-unhighlight", { name: d.node.name });
     })
     .on("click", d => {
       // Stop this event's propagation so that it isn't also considered
       // a graph-select.
       (<Event>d3.event).stopPropagation();
-      sceneBehavior.fire("node-select", { name: d.node.name });
+      sceneElement.fire("node-select", { name: d.node.name });
     })
     .on("contextmenu", (d, i) => {
-      sceneBehavior.fire("node-select", { name: d.node.name });
+      sceneElement.fire("node-select", { name: d.node.name });
       contextMenuFunction.call(d, i);
     });
 };
@@ -266,13 +266,13 @@ function addInteraction(selection, d: render.RenderNodeInfo,
 /**
  * Returns the d3 context menu specification for the provided node.
  */
-export function getContextMenu(node: Node, sceneBehavior) {
+export function getContextMenu(node: Node, sceneElement) {
   let menu = [{
     title: d => {
-      return tf.graph.getIncludeNodeButtonString(node.include);
+      return getIncludeNodeButtonString(node.include);
     },
     action: (elm, d, i) => {
-      sceneBehavior.fire("node-toggle-extract", { name: node.name });
+      sceneElement.fire("node-toggle-extract", { name: node.name });
     }
   }];
   if (canBeInSeries(node)) {
@@ -281,7 +281,7 @@ export function getContextMenu(node: Node, sceneBehavior) {
         return getGroupSettingLabel(node);
       },
       action: (elm, d, i) => {
-        sceneBehavior.fire("node-toggle-seriesgroup",
+        sceneElement.fire("node-toggle-seriesgroup",
           { name: getSeriesName(node) });
       }
     });
@@ -343,10 +343,10 @@ export function getGroupSettingLabel(node: Node) {
  * Append svg text for label and assign data.
  * @param nodeGroup
  * @param renderNodeInfo The render node information for the label.
- * @param sceneBehavior parent scene module.
+ * @param sceneElement <tf-graph-scene> polymer element.
  */
 function labelBuild(nodeGroup, renderNodeInfo: render.RenderNodeInfo,
-    sceneBehavior) {
+    sceneElement) {
   let namePath = renderNodeInfo.node.name.split("/");
   let text = namePath[namePath.length - 1];
 
@@ -355,13 +355,18 @@ function labelBuild(nodeGroup, renderNodeInfo: render.RenderNodeInfo,
     !renderNodeInfo.expanded;
 
   let label = scene.selectOrCreateChild(nodeGroup, "text", Class.Node.LABEL);
+
+  // Make sure the label is visually on top among its siblings.
+  let labelNode = <HTMLElement> label.node();
+  labelNode.parentNode.appendChild(labelNode);
+
   label.attr("dy", ".35em")
     .attr("text-anchor", "middle");
   if (useFontScale) {
-    if (text.length > sceneBehavior.maxMetanodeLabelLength) {
-      text = text.substr(0, sceneBehavior.maxMetanodeLabelLength - 2) + "...";
+    if (text.length > sceneElement.maxMetanodeLabelLength) {
+      text = text.substr(0, sceneElement.maxMetanodeLabelLength - 2) + "...";
     }
-    let scale = getLabelFontScale(sceneBehavior);
+    let scale = getLabelFontScale(sceneElement);
     label.attr("font-size", scale(text.length) + "px");
   }
   label.text(text);
@@ -373,13 +378,13 @@ function labelBuild(nodeGroup, renderNodeInfo: render.RenderNodeInfo,
  * initialized once by getLabelFontScale.
  */
 let fontScale = null;
-function getLabelFontScale(sceneBehavior) {
+function getLabelFontScale(sceneElement) {
   if (!fontScale) {
     fontScale = d3.scale.linear()
-      .domain([sceneBehavior.maxMetanodeLabelLengthLargeFont,
-        sceneBehavior.maxMetanodeLabelLength])
-      .range([sceneBehavior.maxMetanodeLabelLengthFontSize,
-        sceneBehavior.minMetanodeLabelLengthFontSize]).clamp(true);
+      .domain([sceneElement.maxMetanodeLabelLengthLargeFont,
+        sceneElement.maxMetanodeLabelLength])
+      .range([sceneElement.maxMetanodeLabelLengthFontSize,
+        sceneElement.minMetanodeLabelLengthFontSize]).clamp(true);
   }
   return fontScale;
 }
@@ -401,13 +406,11 @@ function labelPosition(nodeGroup, cx: number, cy: number,
  * @param nodeGroup
  * @param d Render node information.
  * @param nodeClass class for the element.
- * @param before Reference DOM node for insertion.
  * @return Selection of the shape.
  */
-export function buildShape(nodeGroup, d, nodeClass: string, before?) {
+export function buildShape(nodeGroup, d, nodeClass: string) {
   // Create a group to house the underlying visual elements.
-  let shapeGroup = scene.selectOrCreateChild(nodeGroup, "g", nodeClass,
-    before);
+  let shapeGroup = scene.selectOrCreateChild(nodeGroup, "g", nodeClass);
   // TODO(jimbo): DOM structure should be templated in HTML somewhere, not JS.
   switch (d.node.type) {
     case NodeType.OP:
@@ -458,7 +461,7 @@ export function nodeClass(d: render.RenderNodeInfo) {
 };
 
 /** Modify node and its subscene and its label's positional attributes */
-function position(nodeGroup, d: render.RenderNodeInfo, sceneBehavior) {
+function position(nodeGroup, d: render.RenderNodeInfo) {
   let shapeGroup = scene.selectChild(nodeGroup, "g", Class.Node.SHAPE);
   let cx = layout.computeCXPositionOfNodeShape(d);
   switch (d.node.type) {
@@ -520,15 +523,15 @@ export enum ColorBy { STRUCTURE, DEVICE, COMPUTE_TIME, MEMORY };
  */
 export function getFillForNode(templateIndex, colorBy,
     renderInfo: render.RenderNodeInfo, isExpanded: boolean): string {
-  let colorParams = tf.graph.render.MetanodeColors;
+  let colorParams = render.MetanodeColors;
   switch (colorBy) {
     case ColorBy.STRUCTURE:
-      if (renderInfo.node.type === tf.graph.NodeType.META) {
+      if (renderInfo.node.type === NodeType.META) {
         let tid = (<Metanode>renderInfo.node).templateId;
         return tid === null ?
           colorParams.UNKNOWN :
           colorParams.STRUCTURE_PALETTE(templateIndex(tid), isExpanded);
-      } else if (renderInfo.node.type === tf.graph.NodeType.SERIES) {
+      } else if (renderInfo.node.type === NodeType.SERIES) {
         // If expanded, we're showing the background rect, which we want to
         // appear gray. Otherwise we're showing a stack of ellipses which we
         // want to show white.
@@ -587,10 +590,10 @@ export function getFillForNode(templateIndex, colorBy,
  * that can't be done in css).
  */
 export function stylize(nodeGroup, renderInfo: render.RenderNodeInfo,
-    sceneBehavior, nodeClass?) {
+    sceneElement, nodeClass?) {
   nodeClass = nodeClass || Class.Node.SHAPE;
-  let isHighlighted = sceneBehavior.isNodeHighlighted(renderInfo.node.name);
-  let isSelected = sceneBehavior.isNodeSelected(renderInfo.node.name);
+  let isHighlighted = sceneElement.isNodeHighlighted(renderInfo.node.name);
+  let isSelected = sceneElement.isNodeSelected(renderInfo.node.name);
   let isExtract = renderInfo.isInExtract || renderInfo.isOutExtract;
   let isExpanded = renderInfo.expanded;
   nodeGroup.classed("highlighted", isHighlighted);
@@ -601,8 +604,8 @@ export function stylize(nodeGroup, renderInfo: render.RenderNodeInfo,
   // Main node always exists here and it will be reached before subscene,
   // so d3 selection is fine here.
   let node = nodeGroup.select("." + nodeClass + " ." + Class.Node.COLOR_TARGET);
-  let fillColor = getFillForNode(sceneBehavior.templateIndex,
-    ColorBy[sceneBehavior.colorBy.toUpperCase()],
+  let fillColor = getFillForNode(sceneElement.templateIndex,
+    ColorBy[sceneElement.colorBy.toUpperCase()],
     renderInfo, isExpanded);
   node.style("fill", fillColor);
 
@@ -617,7 +620,7 @@ export function stylize(nodeGroup, renderInfo: render.RenderNodeInfo,
 export function getStrokeForFill(fill: string) {
   // If node is colored by a gradient, then use a dark gray outline.
   return fill.substring(0, 3) === "url" ?
-    tf.graph.render.MetanodeColors.GRADIENT_OUTLINE :
+    render.MetanodeColors.GRADIENT_OUTLINE :
     d3.rgb(fill).darker().toString();
 }
 
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/scene.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/scene.ts
index 24c16e31ee..685ad646f7 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/scene/scene.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/scene/scene.ts
@@ -243,12 +243,12 @@ export function selectChild(container, tagName: string, className?: string) {
  *
  * @param container D3 selection of the parent.
  * @param renderNode render node of a metanode or series node.
- * @param sceneBehavior Parent scene module.
+ * @param sceneElement <tf-graph-scene> polymer element.
  * @param sceneClass class attribute of the scene (default="scene").
  */
 export function buildGroup(container,
     renderNode: render.RenderGroupNodeInfo,
-    sceneBehavior,
+    sceneElement,
     sceneClass: string) {
   sceneClass = sceneClass || Class.Scene.GROUP;
   let isNewSceneGroup = selectChild(container, "g", sceneClass).empty();
@@ -272,17 +272,17 @@ export function buildGroup(container,
   }
 
   // Create the layer of edges for this scene (paths).
-  edge.buildGroup(coreGroup, renderNode.coreGraph, sceneBehavior);
+  edge.buildGroup(coreGroup, renderNode.coreGraph, sceneElement);
 
   // Create the layer of nodes for this scene (ellipses, rects etc).
-  node.buildGroup(coreGroup, coreNodes, sceneBehavior);
+  node.buildGroup(coreGroup, coreNodes, sceneElement);
 
   // In-extract
   if (renderNode.isolatedInExtract.length > 0) {
     let inExtractGroup = selectOrCreateChild(sceneGroup, "g",
       Class.Scene.INEXTRACT);
     node.buildGroup(inExtractGroup, renderNode.isolatedInExtract,
-        sceneBehavior);
+        sceneElement);
   } else {
     selectChild(sceneGroup, "g", Class.Scene.INEXTRACT).remove();
   }
@@ -292,7 +292,7 @@ export function buildGroup(container,
     let outExtractGroup = selectOrCreateChild(sceneGroup, "g",
       Class.Scene.OUTEXTRACT);
     node.buildGroup(outExtractGroup, renderNode.isolatedOutExtract,
-        sceneBehavior);
+        sceneElement);
   } else {
     selectChild(sceneGroup, "g", Class.Scene.OUTEXTRACT).remove();
   }
@@ -345,9 +345,9 @@ function position(sceneGroup, renderNode: render.RenderGroupNodeInfo) {
 };
 
 /** Adds a click listener to a group that fires a graph-select event */
-export function addGraphClickListener(graphGroup, sceneBehavior) {
+export function addGraphClickListener(graphGroup, sceneElement) {
   d3.select(graphGroup).on("click", () => {
-    sceneBehavior.fire("graph-select");
+    sceneElement.fire("graph-select");
   });
 };