Adding stats (run metadata) to the Graph UI.

Todos:
- Add actual memory and cpu data to the info card.
Change: 119050200

14 files changed, 372 insertions, 251 deletions

diff --git a/tensorflow/tensorboard/components/tf-graph-app/tf-graph-app.html b/tensorflow/tensorboard/components/tf-graph-app/tf-graph-app.html
index a21202b834..a98a1e155c 100644
--- a/tensorflow/tensorboard/components/tf-graph-app/tf-graph-app.html
+++ b/tensorflow/tensorboard/components/tf-graph-app/tf-graph-app.html
@@ -45,14 +45,13 @@ Example
   <div class="side">
     <tf-graph-controls
         color-by-params="[[colorByParams]]"
-        has-stats="[[hasStats]]"
+        stats="[[stats]]"
         color-by="{{colorBy}}"
     ></tf-graph-controls>
     <tf-graph-loader id="loader"
         out-graph-hierarchy="{{graphHierarchy}}"
         out-graph="{{graph}}"
-        out-graph-name="{{graphName}}"
-        has-stats="{{hasStats}}"
+        out-stats="{{stats}}"
         progress="{{_progress}}"
     ></tf-graph-loader>
   </div>
@@ -60,8 +59,7 @@ Example
     <tf-graph-board id="graphboard"
         graph-hierarchy="[[graphHierarchy]]"
         graph="[[graph]]"
-        has-stats="[[hasStats]]"
-        graph-name="[[graphName]]"
+        stats="[[stats]]"
         progress="[[_progress]]"
         color-by="[[colorBy]]"
         color-by-params="{{colorByParams}}"
@@ -77,7 +75,7 @@ Example
 Polymer({
   is: 'tf-graph-app',
   properties: {
-    hasStats: Boolean,
+    stats: Object,
     pbtxt: {
       type: String,
       observer: '_updateGraph',
diff --git a/tensorflow/tensorboard/components/tf-graph-board/tf-graph-board.html b/tensorflow/tensorboard/components/tf-graph-board/tf-graph-board.html
index 4f00686566..e352038ea6 100644
--- a/tensorflow/tensorboard/components/tf-graph-board/tf-graph-board.html
+++ b/tensorflow/tensorboard/components/tf-graph-board/tf-graph-board.html
@@ -117,11 +117,11 @@ paper-progress {
               basic-graph="[[graph]]"
               hierarchy-params="[[hierarchyParams]]"
               render-hierarchy="{{_renderHierarchy}}"
+              stats="[[stats]]"
               selected-node="{{_selectedNode}}"
               highlighted-node="{{_highlightedNode}}"
               color-by="[[colorBy]]"
               color-by-params="{{colorByParams}}"
-              graph-name="[[graphName]]"
               progress="{{progress}}"
     ></tf-graph>
   </div>
@@ -150,9 +150,7 @@ Polymer({
     // Public API.
     graphHierarchy: Object,
     graph: Object,
-    graphName: String,
-    // True if the graph data has also run-time stats.
-    hasStats: Boolean,
+    stats: Object,
     /**
      * @type {value: number, msg: string}
      *
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/common.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/common.ts
index ab301a210e..18c35f1d96 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/common.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/common.ts
@@ -239,31 +239,31 @@ export interface TFNode {
 /**
  * TensorFlow stats file definition as defined in the stats proto file.
  */
-export interface TFStats {
-  devStats: {device: string, nodeStats: TFNodeStats[]}[];
+export interface StepStats {
+  dev_stats: {device: string, node_stats: NodeStats[]}[];
 }
 
 /**
  * TensorFlow stats for a node as defined in the stats proto file.
  */
-export interface TFNodeStats {
-  nodeName: string;
+export interface NodeStats {
+  node_name: string;
   // The next 4 properties are currently stored as string in json
   // and must be parsed.
-  allStartMicros: number;
-  opStartRelMicros: number;
-  opEndRelMicros: number;
-  allEndRelMicros: number;
+  all_start_micros: number;
+  op_start_rel_micros: number;
+  op_end_rel_micros: number;
+  all_end_rel_micros: number;
   memory: {
-    allocatorName: string;
-    totalBytes: number; // Stored as string in json and should be parsed.
-    peakBytes: number; // Stored as string in json and should be parsed.
+    allocator_name: string;
+    total_bytes: number; // Stored as string in json and should be parsed.
+    peak_bytes: number; // Stored as string in json and should be parsed.
   }[];
   /** Output sizes recorded for a single execution of a graph node */
   output: TFNodeOutput[];
-  timelineLabel: string;
-  scheduledMicros: string;
-  threadId: string;
+  timeline_label: string;
+  scheduled_micros: string;
+  thread_id: string;
 }
 
 /**
@@ -271,9 +271,7 @@ export interface TFNodeStats {
  */
 export interface TFNodeOutput {
   slot: number; // Stored as string in json and should be parsed.
-  /** Was the tensor allocated by this Op or a previous computation */
-  allocationType: string;
-  tensorDescription: {
+  tensor_description: {
     /** Data type of tensor elements */
     dtype: string;
     /** Shape of the tensor */
@@ -292,15 +290,15 @@ export interface TFNodeOutput {
       }[];
     };
     /** Information about the size and allocator used for the data */
-    allocationDescription: {
+    allocation_description: {
       // The next 2 properties are stored as string in json and
       // should be parsed.
       /** Total number of bytes requested */
-      requestedBytes: number;
+      requested_bytes: number;
       /** Total number of bytes allocated, if known */
-      allocatedBytes?: number;
+      allocated_bytes?: number;
       /** Name of the allocator used */
-      allocatorName: string;
+      allocator_name: string;
     };
   };
 }
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/graph.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/graph.ts
index b2f6d21598..c5266565bd 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/graph.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/graph.ts
@@ -376,35 +376,40 @@ export function createMetanode(name: string, opt = {}): Metanode {
  * graph information.
  */
 export function joinStatsInfoWithGraph(graph: SlimGraph,
-    statsJson: TFStats): void {
-  _.each(statsJson.devStats, stats => {
-    _.each(stats.nodeStats, nodeStats => {
+    stats: StepStats): void {
+  _.each(stats.dev_stats, devStats => {
+    _.each(devStats.node_stats, nodeStats => {
       // Lookup the node in the graph by its original name, e.g. A. If not
       // found, lookup by the rewritten name A/(A) in case the name is both
       // a namespace and a node name.
-      let nodeName = nodeStats.nodeName in graph.nodes ?
-          nodeStats.nodeName :
-          nodeStats.nodeName + NAMESPACE_DELIM + "(" + nodeStats.nodeName + ")";
-      if (nodeName in graph.nodes) {
-        // Compute the total bytes used.
-        let totalBytes = 0;
-        if (nodeStats.memory) {
-          _.each(nodeStats.memory, alloc => {
-            if (alloc.totalBytes) {
-              totalBytes += Number(alloc.totalBytes);
-            }
-          });
-        }
-        let outputSize: number[][] = null;
-        if (nodeStats.output) {
-          outputSize = _.map(nodeStats.output, output => {
-            return _.map(output.tensorDescription.shape.dim,
-                dim => Number(dim.size));
-          });
-        }
-        graph.nodes[nodeName].stats = new NodeStats(totalBytes,
-            Number(nodeStats.allEndRelMicros), outputSize);
+      let nodeName = nodeStats.node_name in graph.nodes ?
+          nodeStats.node_name :
+          nodeStats.node_name + NAMESPACE_DELIM + "(" + nodeStats.node_name + ")";
+
+      // Couldn't find a matching node.
+      if (!(nodeName in graph.nodes)) {
+        return;
+      }
+
+      // Compute the total bytes used.
+      let totalBytes = 0;
+      if (nodeStats.memory) {
+        _.each(nodeStats.memory, alloc => {
+          if (alloc.total_bytes) {
+            totalBytes += Number(alloc.total_bytes);
+          }
+        });
+      }
+      let outputSize: number[][] = null;
+      if (nodeStats.output) {
+        outputSize = _.map(nodeStats.output, output => {
+          return _.map(output.tensor_description.shape.dim,
+              dim => Number(dim.size));
+        });
       }
+      graph.nodes[nodeName].device = devStats.device;
+      graph.nodes[nodeName].stats = new NodeStats(totalBytes,
+          Number(nodeStats.all_end_rel_micros), outputSize);
     });
   });
 }
@@ -492,7 +497,6 @@ class MetanodeImpl implements Metanode {
     this.templateId = null;
     /** Metanode which contains this node, if any */
     this.parentNode = null;
-    this.stats = new NodeStats(0, 0, null);
     this.hasNonControlEdges = false;
     this.include = InclusionType.UNSPECIFIED;
   }
@@ -705,7 +709,6 @@ class SeriesNodeImpl implements SeriesNode {
     this.parentNode = null;
     this.deviceHistogram = {};
     this.hasNonControlEdges = false;
-    this.stats = new NodeStats(0, 0, null);
     this.include = InclusionType.UNSPECIFIED;
   }
 }
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/hierarchy.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/hierarchy.ts
index 8dca63c9ab..ef6d69d4d5 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/hierarchy.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/hierarchy.ts
@@ -428,6 +428,42 @@ export function build(graph: tf.graph.SlimGraph, params: HierarchyParams,
   });
 };
 
+export function joinAndAggregateStats(h: Hierarchy, stats: StepStats) {
+  // Get all the possible device names.
+  let deviceNames = {};
+  _.each(h.root.leaves(), nodeName => {
+    let leaf = <OpNode> h.node(nodeName);
+    if (leaf.device != null) {
+      deviceNames[leaf.device] = true;
+    }
+  });
+  h.devices = _.keys(deviceNames);
+
+  // Reset stats for each group node.
+  _.each(h.getNodeMap(), (node, nodeName) => {
+    if (node.isGroupNode) {
+      node.stats = new NodeStats(0, 0, null);
+      (<GroupNode>node).deviceHistogram = {};
+    }
+  });
+
+  // Bubble-up the stats and device distribution from leaves to parents.
+  _.each(h.root.leaves(), nodeName => {
+    let leaf = <OpNode> h.node(nodeName);
+    let node = <GroupNode|OpNode> leaf;
+    while (node.parentNode != null) {
+      if (leaf.device != null) {
+        let deviceHistogram = (<GroupNode>node.parentNode).deviceHistogram;
+        deviceHistogram[leaf.device] = (deviceHistogram[leaf.device] || 0) + 1;
+      }
+      if (leaf.stats != null) {
+        node.parentNode.stats.combine(leaf.stats);
+      }
+      node = <GroupNode> node.parentNode;
+    }
+  });
+}
+
 /**
  * Creates the metanodes in the hierarchical graph and assigns parent-child
  * relationship between them.
@@ -446,9 +482,6 @@ function addNodes(h: Hierarchy, graph: SlimGraph) {
       parent.depth = Math.max(parent.depth, path.length - i);
       parent.cardinality += node.cardinality;
       parent.opHistogram[node.op] = (parent.opHistogram[node.op] || 0) + 1;
-      if (node.stats) {
-        parent.stats.combine(node.stats);
-      }
       if (node.device != null) {
         parent.deviceHistogram[node.device] =
             (parent.deviceHistogram[node.device] || 0) + 1;
@@ -623,11 +656,6 @@ function groupSeries(metanode: Metanode, hierarchy: Hierarchy,
       }
       child.parentNode = seriesNode;
       seriesNames[n] = seriesName;
-
-      if (child.stats) {
-        seriesNode.stats.combine(child.stats);
-      }
-
       // Remove now-grouped node from its original parent's metagraph.
       metagraph.removeNode(n);
     });
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/parser.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/parser.ts
index 0ed1bd1961..865b6e6761 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/parser.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/parser.ts
@@ -37,7 +37,7 @@ function parseValue(value: string): string|number|boolean {
 /**
  * Fetches a text file and returns a promise of the result.
  */
-export function readPbTxt(filepath: string): Promise<string> {
+export function fetchPbTxt(filepath: string): Promise<string> {
   return new Promise<string>(function(resolve, reject) {
     d3.text(filepath, function(error, text) {
       if (error) {
@@ -50,52 +50,36 @@ export function readPbTxt(filepath: string): Promise<string> {
 }
 
 /**
- * Fetches and parses a json file and returns a promise of the result.
+ * Fetches the metadata file, parses it and returns a promise of the result.
  */
-export function readJson(filepath: string): Promise<Object> {
-  return new Promise<Object>(function(resolve, reject) {
-    d3.json(filepath, function(error, text) {
-      if (error) {
-        reject(error);
-        return;
-      }
-      resolve(text);
-    });
+export function fetchAndParseMetadata(path: string, tracker: ProgressTracker) {
+  return runTask("Reading metadata pbtxt", 40, () => {
+    if (path == null) {
+      return Promise.resolve(null);
+    }
+    return fetchPbTxt(path).then(text => new Blob([text]));
+  }, tracker)
+  .then((blob: Blob) => {
+    return runTask("Parsing metadata.pbtxt", 60, () => {
+      return blob != null ? parseStatsPbTxt(blob) : null;
+    }, tracker);
   });
 }
 
 /**
- * Reads the graph and stats file (if available), parses them and returns a
- * promise of the result.
+ * Fetches the graph file, parses it and returns a promise of the result.
  */
-export function readAndParseData(dataset: {path: string, statsPath: string},
-    pbTxtFile: Blob, tracker: ProgressTracker):
-    Promise<{ nodes: TFNode[], statsJson: Object }|void> {
-  let graphPbTxt: Blob;
-  let statsJson: Object;
-  return runTask("Reading graph.pbtxt", 20, () => {
+export function fetchAndParseGraphData(path: string, pbTxtFile: Blob,
+    tracker: ProgressTracker) {
+  return runTask("Reading graph pbtxt", 40, () => {
     return pbTxtFile ?
       Promise.resolve(pbTxtFile) :
-      readPbTxt(dataset.path).then(text => new Blob([text]));
+      fetchPbTxt(path).then(text => new Blob([text]));
   }, tracker)
   .then(blob => {
-    graphPbTxt = blob;
-    return runTask("Reading stats.pbtxt", 20, () => {
-      return (dataset != null && dataset.statsPath != null) ?
-          readJson(dataset.statsPath) : null;
-    }, tracker);
-  })
-  .then(json => {
-    statsJson = json;
     return runTask("Parsing graph.pbtxt", 60, () => {
-      return parsePbtxtFile(graphPbTxt);
+      return parseGraphPbTxt(blob);
     }, tracker);
-  })
-  .then(nodes => {
-    return {
-      nodes: nodes,
-      statsJson: statsJson
-    };
   });
 }
 
@@ -158,13 +142,59 @@ export function streamParse(file: Blob, callback: (string) => void,
 }
 
 /**
- * Parses a proto txt file or blob into javascript object.
+ * Since proto-txt doesn't explicitly say whether an attribute is repeated
+ * (an array) or not, we keep a hard-coded list of attributes that are known
+ * to be repeated. This list is used in parsing time to convert repeated
+ * attributes into arrays even when the attribute only shows up once in the
+ * object.
+ */
+const GRAPH_REPEATED_FIELDS: {[attrPath: string]: boolean} = {
+  "node": true,
+  "node.input": true,
+  "node.attr": true,
+  "node.attr.value.list.type": true,
+  "node.attr.value.shape.dim": true,
+  "node.attr.value.tensor.string_val": true,
+  "node.attr.value.tensor.tensor_shape.dim": true,
+  "node.attr.value.list.shape": true,
+  "node.attr.value.list.shape.dim": true,
+  "node.attr.value.list.s": true
+};
+
+const METADATA_REPEATED_FIELDS: {[attrPath: string]: boolean} = {
+  "step_stats.dev_stats": true,
+  "step_stats.dev_stats.node_stats": true,
+  "step_stats.dev_stats.node_stats.output": true,
+  "step_stats.dev_stats.node_stats.memory": true,
+  "step_stats.dev_stats.node_stats.output.tensor_description.shape.dim": true
+};
+
+/**
+ * Parses a blob of proto txt file into a raw Graph object.
+ */
+export function parseGraphPbTxt(input: Blob): Promise<TFNode[]> {
+  return parsePbtxtFile(input, GRAPH_REPEATED_FIELDS).then(obj => obj["node"]);
+}
+
+/**
+ * Parses a blob of proto txt file into a StepStats object.
+ */
+function parseStatsPbTxt(input: Blob): Promise<StepStats> {
+  return parsePbtxtFile(input, METADATA_REPEATED_FIELDS)
+    .then(obj => obj["step_stats"]);
+}
+
+/**
+ * Parses a blob of proto txt file into javascript object.
  *
  * @param input The Blob or file object implementing slice.
+ * @param repeatedFields Map (Set) of all the repeated fields, since you can't
+ *   tell directly from the pbtxt if a field is repeated or not.
  * @returns The parsed object.
  */
-export function parsePbtxtFile(input: Blob): Promise<TFNode[]> {
-  let output: { [name: string]: any; } = { node: [] };
+function parsePbtxtFile(input: Blob,
+    repeatedFields: {[attrPath: string]: boolean}): Promise<Object> {
+  let output: { [name: string]: any; } = {};
   let stack = [];
   let path: string[] = [];
   let current: { [name: string]: any; } = output;
@@ -180,26 +210,6 @@ export function parsePbtxtFile(input: Blob): Promise<TFNode[]> {
   }
 
   /**
-   * Since proto-txt doesn't explicitly say whether an attribute is repeated
-   * (an array) or not, we keep a hard-coded list of attributes that are known
-   * to be repeated. This list is used in parsing time to convert repeated
-   * attributes into arrays even when the attribute only shows up once in the
-   * object.
-   */
-  let ARRAY_ATTRIBUTES: {[attrPath: string]: boolean} = {
-    "node": true,
-    "node.input": true,
-    "node.attr": true,
-    "node.attr.value.list.type": true,
-    "node.attr.value.shape.dim": true,
-    "node.attr.value.tensor.string_val": true,
-    "node.attr.value.tensor.tensor_shape.dim": true,
-    "node.attr.value.list.shape": true,
-    "node.attr.value.list.shape.dim": true,
-    "node.attr.value.list.s": true
-  };
-
-  /**
    * Adds a value, given the attribute name and the host object. If the
    * attribute already exists, but is not an array, it will convert it to an
    * array of values.
@@ -215,7 +225,7 @@ export function parsePbtxtFile(input: Blob): Promise<TFNode[]> {
     // We treat "node" specially since it is done so often.
     let existingValue = obj[name];
     if (existingValue == null) {
-      obj[name] = path.join(".") in ARRAY_ATTRIBUTES ? [value] : value;
+      obj[name] = path.join(".") in repeatedFields ? [value] : value;
     } else if (Array.isArray(existingValue)) {
       existingValue.push(value);
     } else {
@@ -247,19 +257,8 @@ export function parsePbtxtFile(input: Blob): Promise<TFNode[]> {
         break;
     }
   }).then(function() {
-    return output["node"];
+    return output;
   });
 }
 
-/**
- * Parses a proto txt file into a javascript object.
- *
- * @param input The string contents of the proto txt file.
- * @return The parsed object.
- */
-export function parsePbtxt(input: string): Promise<TFNode[]> {
-  let blob = new Blob([input]);
-  return parsePbtxtFile(blob);
-}
-
 } // Close module tf.graph.parser.
diff --git a/tensorflow/tensorboard/components/tf-graph-common/lib/render.ts b/tensorflow/tensorboard/components/tf-graph-common/lib/render.ts
index 5110e2f3f7..4033219b62 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/lib/render.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/lib/render.ts
@@ -167,12 +167,24 @@ export class RenderGraphInfo {
   constructor(hierarchy: hierarchy.Hierarchy) {
     this.hierarchy = hierarchy;
     this.index = {};
+
+    this.computeScales();
+    // Maps node name to whether the rendering hierarchy was already
+    // constructed.
+    this.hasSubhierarchy = {};
+    this.root = new RenderGroupNodeInfo(hierarchy.root);
+    this.index[hierarchy.root.name] = this.root;
+    this.buildSubhierarchy(hierarchy.root.name);
+    this.root.expanded = true;
+  }
+
+  computeScales() {
     this.deviceColorMap = d3.scale.ordinal<string>()
-        .domain(hierarchy.devices)
-        .range(_.map(d3.range(hierarchy.devices.length),
+        .domain(this.hierarchy.devices)
+        .range(_.map(d3.range(this.hierarchy.devices.length),
                      MetanodeColors.DEVICE_PALETTE));
 
-    let topLevelGraph = hierarchy.root.metagraph;
+    let topLevelGraph = this.hierarchy.root.metagraph;
     // Find the maximum and minimum memory usage.
     let memoryExtent = d3.extent(topLevelGraph.nodes(),
         (nodeName, index) => {
@@ -198,14 +210,6 @@ export class RenderGraphInfo {
     this.computeTimeScale = d3.scale.linear<string, string>()
         .domain(computeTimeExtent)
         .range(PARAMS.minMaxColors);
-
-    // Maps node name to whether the rendering hierarchy was already
-    // constructed.
-    this.hasSubhierarchy = {};
-    this.root = new RenderGroupNodeInfo(hierarchy.root);
-    this.index[hierarchy.root.name] = this.root;
-    this.buildSubhierarchy(hierarchy.root.name);
-    this.root.expanded = true;
   }
 
   /**
diff --git a/tensorflow/tensorboard/components/tf-graph-common/test/parser-test.ts b/tensorflow/tensorboard/components/tf-graph-common/test/parser-test.ts
index cc4c951f7d..7d510c57ae 100644
--- a/tensorflow/tensorboard/components/tf-graph-common/test/parser-test.ts
+++ b/tensorflow/tensorboard/components/tf-graph-common/test/parser-test.ts
@@ -32,7 +32,7 @@ test("simple pbtxt", (done) => {
        input: "Q"
        input: "W"
      }`;
-  tf.graph.parser.parsePbtxt(pbtxt).then(nodes => {
+  tf.graph.parser.parseGraphPbTxt(new Blob([pbtxt])).then(nodes => {
     assert.isTrue(nodes != null && nodes.length === 3);
     done();
   });
diff --git a/tensorflow/tensorboard/components/tf-graph-dashboard/tf-graph-dashboard.html b/tensorflow/tensorboard/components/tf-graph-dashboard/tf-graph-dashboard.html
index 616e701973..3fb87417e9 100644
--- a/tensorflow/tensorboard/components/tf-graph-dashboard/tf-graph-dashboard.html
+++ b/tensorflow/tensorboard/components/tf-graph-dashboard/tf-graph-dashboard.html
@@ -27,20 +27,21 @@ by default. The user can select a different run from a dropdown menu.
 <div class="sidebar">
   <tf-graph-controls id="controls"
       color-by-params="[[_colorByParams]]"
-      has-stats="[[_hasStats]]"
+      stats="[[_stats]]"
       color-by="{{_colorBy}}",
       datasets="[[_datasets]]",
       selected-dataset="{{_selectedDataset}}"
       selected-file="{{_selectedFile}}"
+      selected-metadata-tag="{{_selectedMetadataTag}}"
   ></tf-graph-controls>
   <tf-graph-loader id="loader"
           datasets="[[_datasets]]",
           selected-dataset="[[_selectedDataset]]"
+          selected-metadata-tag="[[_selectedMetadataTag]]"
           selected-file="[[_selectedFile]]"
           out-graph-hierarchy="{{_graphHierarchy}}"
           out-graph="{{_graph}}"
-          out-graph-name="{{_graphName}}"
-          has-stats="{{_hasStats}}"
+          out-stats="{{_stats}}"
           progress="{{_progress}}"
           out-hierarchy-params="{{_hierarchyParams}}"
   ></tf-graph-loader>
@@ -49,8 +50,7 @@ by default. The user can select a different run from a dropdown menu.
     <tf-graph-board id="graphboard"
                 graph-hierarchy="[[_graphHierarchy]]"
                 graph="[[_graph]]"
-                has-stats="[[_hasStats]]"
-                graph-name="[[_graphName]]"
+                stats="[[_stats]]"
                 progress="[[_progress]]"
                 color-by="[[_colorBy]]"
                 color-by-params="{{_colorByParams}}"
@@ -79,28 +79,31 @@ by default. The user can select a different run from a dropdown menu.
 Polymer({
   is: 'tf-graph-dashboard',
   properties: {
-    _datasets: {
-      type: Object,
-      computed: '_getDatasets(runs.*, router)'
-    },
+    _datasets: Object,
     backend: {type: Object, observer: 'reload'},
     router: {type: Object},
     runs: Array,
   },
   reload: function() {
-    var _this = this;
-    this.backend.graphRuns().then(function(x) {
-      _this.runs = x;
-    });
-  },
-  _getDatasets: function(runs, router) {
-    return _.map(this.runs, function(runName) {
-      return {
-        name: runName,
-        path: router.graph(runName, tf.graph.LIMIT_ATTR_SIZE,
-            tf.graph.LARGE_ATTRS_KEY)
-      };
-    });
+    Promise.all([this.backend.graphRuns(), this.backend.runMetadataRuns()])
+    .then(function(result) {
+      var runsWithGraph = result[0];
+      var runToMetadata = result[1];
+      var datasets = _.map(runsWithGraph, function(runName) {
+        return {
+          name: runName,
+          path: this.router.graph(runName, tf.graph.LIMIT_ATTR_SIZE,
+            tf.graph.LARGE_ATTRS_KEY),
+          runMetadata: _.map(runToMetadata[runName], function(tag) {
+            return {
+              tag: tag,
+              path: this.router.runMetadata(tag, runName)
+            };
+          }, this)
+        };
+      }, this);
+      this.set('_datasets', datasets);
+    }.bind(this));
   },
   _datasetsEmpty: function(datasets) {
     return !datasets || !datasets.length;
diff --git a/tensorflow/tensorboard/components/tf-graph-loader/tf-graph-loader.html b/tensorflow/tensorboard/components/tf-graph-loader/tf-graph-loader.html
index cba3dfa13a..5f9dc847be 100644
--- a/tensorflow/tensorboard/components/tf-graph-loader/tf-graph-loader.html
+++ b/tensorflow/tensorboard/components/tf-graph-loader/tf-graph-loader.html
@@ -23,11 +23,6 @@ Polymer({
       notify: true,
     },
     datasets: Array,
-    hasStats: {
-      type: Boolean,
-      readOnly: true, // This property produces data.
-      notify: true
-    },
     selectedDataset: Number,
     selectedFile: {
       type: Object,
@@ -43,21 +38,41 @@ Polymer({
       readOnly: true, //readonly so outsider can't change this via binding
       notify: true
     },
-    outGraphName: {
-      type: String,
-      readOnly: true,
-      notify: true
-    },
     outHierarchyParams: {
       type: Object,
       readOnly: true,
       notify: true
     },
+    outStats: {
+      type: Object,
+      readOnly: true, // This property produces data.
+      notify: true
+    }
   },
   observers: [
-    '_selectedDatasetChanged(selectedDataset, datasets)'
+    '_selectedDatasetChanged(selectedDataset, datasets)',
+    '_readAndParseMetadata(selectedDataset, selectedMetadataTag, datasets)'
   ],
-  _parseAndConstructHierarchicalGraph: function(dataset, pbTxtFile) {
+  _readAndParseMetadata: function(datasetIndex, metadataIndex, datasets) {
+    if (metadataIndex == -1 || datasets[datasetIndex] == null ||
+        datasets[datasetIndex].runMetadata == null ||
+        datasets[datasetIndex].runMetadata[metadataIndex] == null) {
+      this._setOutStats(null);
+      return;
+    }
+    var path = datasets[datasetIndex].runMetadata[metadataIndex].path;
+    // Reset the progress bar to 0.
+    this.set('progress', {
+      value: 0,
+      msg: ''
+    });
+    var tracker = tf.getTracker(this);
+    tf.graph.parser.fetchAndParseMetadata(path, tracker)
+    .then(function(stats) {
+      this._setOutStats(stats);
+    }.bind(this));
+  },
+  _parseAndConstructHierarchicalGraph: function(path, pbTxtFile) {
     // Reset the progress bar to 0.
     this.set('progress', {
       value: 0,
@@ -76,13 +91,10 @@ Polymer({
       seriesMap: {},
     };
     this._setOutHierarchyParams(hierarchyParams);
-    var statsJson;
     var dataTracker = tf.getSubtaskTracker(tracker, 30, 'Data');
-    tf.graph.parser.readAndParseData(dataset, pbTxtFile, dataTracker)
-    .then(function(result) {
+    tf.graph.parser.fetchAndParseGraphData(path, pbTxtFile, dataTracker)
+    .then(function(graph) {
       // Build the flat graph (consists only of Op nodes).
-      var nodes = result.nodes;
-      statsJson = result.statsJson;
 
       // This is the whitelist of inputs on op types that are considered
       // reference edges. "Assign 0" indicates that the first input to
@@ -107,18 +119,11 @@ Polymer({
         outEmbeddingTypes: ['^[a-zA-Z]+Summary$'],
         refEdges: refEdges
       };
-      var graphTracker = tf.getSubtaskTracker(tracker, 20,
-          'Graph');
-      return tf.graph.build(nodes, buildParams, graphTracker);
+      var graphTracker = tf.getSubtaskTracker(tracker, 20, 'Graph');
+      return tf.graph.build(graph, buildParams, graphTracker);
     })
     .then(function(graph) {
       this._setOutGraph(graph);
-      if (statsJson) {
-        // If there are associated stats, join them with the graph.
-        tf.time('Joining stats info with graph...', function() {
-          tf.graph.joinStatsInfoWithGraph(graph, statsJson);
-        });
-      }
       var hierarchyTracker = tf.getSubtaskTracker(tracker, 50,
           'Namespace hierarchy');
       return tf.graph.hierarchy.build(graph, hierarchyParams, hierarchyTracker);
@@ -126,7 +131,6 @@ Polymer({
     .then(function(graphHierarchy) {
       // Update the properties which notify the parent with the
       // graph hierarchy and whether the data has live stats or not.
-      this._setHasStats(statsJson != null);
       this._setOutGraphHierarchy(graphHierarchy);
     }.bind(this))
     .catch(function(e) {
@@ -136,9 +140,7 @@ Polymer({
     });
   },
   _selectedDatasetChanged: function(datasetIndex, datasets) {
-    var dataset = datasets[datasetIndex];
-    this._parseAndConstructHierarchicalGraph(dataset);
-    this._setOutGraphName(dataset.name);
+    this._parseAndConstructHierarchicalGraph(datasets[datasetIndex].path);
   },
   _selectedFileChanged: function(e) {
     if (!e) {
diff --git a/tensorflow/tensorboard/components/tf-graph/demo/tf-graph-demo.html b/tensorflow/tensorboard/components/tf-graph/demo/tf-graph-demo.html
index 2e5ef3d1de..05256ba399 100644
--- a/tensorflow/tensorboard/components/tf-graph/demo/tf-graph-demo.html
+++ b/tensorflow/tensorboard/components/tf-graph/demo/tf-graph-demo.html
@@ -45,20 +45,21 @@ Example
   <div class="side">
     <tf-graph-controls
         color-by-params="[[colorByParams]]"
-        has-stats="[[hasStats]]"
+        stats="[[stats]]"
         color-by="{{colorBy}}"
         datasets="[[datasets]]",
         selected-dataset="{{selectedDataset}}"
         selected-file="{{selectedFile}}"
+        selected-metadata-tag="{{selectedMetadataTag}}"
     ></tf-graph-controls>
     <tf-graph-loader id="loader"
         datasets="[[datasets]]",
         selected-dataset="[[selectedDataset]]"
+        selected-metadata-tag="[[selectedMetadataTag]]"
         selected-file="[[selectedFile]]"
         out-graph-hierarchy="{{graphHierarchy}}"
         out-graph="{{graph}}"
-        out-graph-name="{{graphName}}"
-        has-stats="{{hasStats}}"
+        out-stats="{{stats}}"
         progress="{{_progress}}"
         out-hierarchy-params="{{_hierarchyParams}}"
     ></tf-graph-loader>
@@ -67,8 +68,7 @@ Example
     <tf-graph-board id="graphboard"
         graph-hierarchy="[[graphHierarchy]]"
         graph="[[graph]]"
-        has-stats="[[hasStats]]"
-        graph-name="[[graphName]]"
+        stats="[[stats]]"
         progress="[[_progress]]"
         color-by="[[colorBy]]"
         color-by-params="{{colorByParams}}"
@@ -88,9 +88,18 @@ var datasets = [
     path: "mnist_eval.pbtxt",
   },
   {
-    name: "Mnist Train (with stats)",
-    path: "mnist_train.pbtxt",
-    statsPath: "mnist_train_stats.json"
+    name: "Mnist with summaries (+stats)",
+    path: "mnist_with_summaries.pbtxt",
+    runMetadata: [
+      {
+        tag: "step100",
+        path: "mnist_with_summaries_step100.pbtxt"
+      },
+      {
+        tag: "step1000",
+        path: "mnist_with_summaries_step1000.pbtxt"
+      }
+    ]
   },
   {
     name: "Mnist Train (with shapes)",
@@ -121,16 +130,48 @@ var datasets = [
     path: "ptb_word_lstm_test_eval.pbtxt",
   },
   {
-    name: "Cifar10 Train (with shapes)",
+    name: "Cifar10 Train (+stats)",
     path: "cifar10_train.pbtxt",
+    runMetadata: [
+      {
+        tag: "step0",
+        path: "cifar10_train_step0.pbtxt"
+      },
+      {
+        tag: "step100",
+        path: "cifar10_train_step100.pbtxt"
+      },
+      {
+        tag: "step200",
+        path: "cifar10_train_step200.pbtxt"
+      },
+      {
+        tag: "step300",
+        path: "cifar10_train_step300.pbtxt"
+      }
+    ]
   },
   {
     name: "Cifar10 Multi-GPU Train",
     path: "cifar10_multi_gpu_train.pbtxt",
   },
   {
-    name: "Cifar10 Eval",
+    name: "Cifar10 Eval (+stats)",
     path: "cifar10_eval.pbtxt",
+    runMetadata: [
+      {
+        tag: "step0",
+        path: "cifar10_eval_step0.pbtxt"
+      },
+      {
+        tag: "step10",
+        path: "cifar10_eval_step10.pbtxt"
+      },
+      {
+        tag: "step20",
+        path: "cifar10_eval_step20.pbtxt"
+      },
+    ]
   },
   {
     name: "Fatcat LSTM",
@@ -161,7 +202,6 @@ var datasets = [
 Polymer({
   is: 'tf-graph-demo',
   properties: {
-    hasStats: Boolean,
     datasets: {
       type: Object,
       value: function() {
@@ -170,24 +210,27 @@ Polymer({
     },
     selectedDataset: {
       type: Number,
-      value: 9 // Cifar train with shapes.
+      value: 1 // Mnist with metadata info.
     },
     _progress: Object
   },
+  _normalizePath: function(path) {
+    return this.resolveUrl('tf_model_zoo/' + path)
+      .replace('tf_model_zoo', DEMO_DIR_PREFIX + 'tf_model_zoo');
+  },
   _getDatasets: function() {
     if(typeof DEMO_DIR_PREFIX === 'undefined') {
       DEMO_DIR_PREFIX = '';
     }
-    return _.map(datasets, function(dataset) {
-      var result = {
-        name: dataset.name,
-        path: this.resolveUrl('tf_model_zoo/' + dataset.path).replace("tf_model_zoo", DEMO_DIR_PREFIX + 'tf_model_zoo')
-      };
-      if (dataset.statsPath != null) {
-        result.statsPath = this.resolveUrl('tf_model_zoo/' + dataset.statsPath).replace("tf_model_zoo", DEMO_DIR_PREFIX + 'tf_model_zoo');
+    _.each(datasets, function(dataset) {
+      dataset.path = this._normalizePath(dataset.path);
+      if (dataset.runMetadata != null) {
+        _.each(dataset.runMetadata, function(metadata) {
+          metadata.path = this._normalizePath(metadata.path);
+        }, this);
       }
-      return result;
     }, this);
+    return datasets;
   }
 });
 })();
diff --git a/tensorflow/tensorboard/components/tf-graph/tf-graph-controls.html b/tensorflow/tensorboard/components/tf-graph/tf-graph-controls.html
index 28b0a607cc..473ca082bc 100644
--- a/tensorflow/tensorboard/components/tf-graph/tf-graph-controls.html
+++ b/tensorflow/tensorboard/components/tf-graph/tf-graph-controls.html
@@ -1,6 +1,7 @@
 <link rel="import" href="../polymer/polymer.html">
 <link rel="import" href="../paper-menu/paper-menu.html">
 <link rel="import" href="../paper-dropdown-menu/paper-dropdown-menu.html">
+<link rel="import" href="../paper-radio-group/paper-radio-group.html">
 
 <dom-module id="tf-graph-controls">
 <template>
@@ -52,6 +53,7 @@ table td {
 }
 
 .allcontrols {
+  width: 188px;
   padding: 30px;
 }
 
@@ -62,6 +64,7 @@ table td {
 }
 
 paper-radio-button {
+  display: block;
   padding: 5px;
 }
 svg.icon {
@@ -125,7 +128,7 @@ svg.icon {
 }
 
 .color-text {
-  padding: 0 0 0 55px;
+  padding: 0 0 0 49px;
 }
 
 .button-text {
@@ -158,6 +161,15 @@ svg.icon {
   display: flex;
   clear: both;
 }
+
+.allcontrols .control-holder paper-radio-group {
+  margin-top: 5px;
+}
+
+span.counter {
+  font-size: 13px;
+  color: gray;
+}
 </style>
 <div class="allcontrols">
   <div class="control-holder">
@@ -175,7 +187,7 @@ svg.icon {
     </a>
   </div>
   <div class="control-holder">
-    <div class="title">Run</div>
+    <div class="title">Run <span class="counter">([[datasets.length]])</span></div>
     <paper-dropdown-menu no-label-float no-animations noink class="run-dropdown">
       <paper-menu id="select" class="dropdown-content" selected="{{selectedDataset}}">
         <template is="dom-repeat" items="[[datasets]]">
@@ -185,6 +197,17 @@ svg.icon {
     </paper-dropdown-menu>
   </div>
   <div class="control-holder">
+    <div class="title">Session runs <span class="counter">([[_numSessionRuns(metadataTags)]])</span></div>
+    <paper-dropdown-menu no-label-float no-animations noink class="run-dropdown">
+      <paper-menu id="select" class="dropdown-content" selected="{{selectedMetadataTag}}">
+        <template is="dom-repeat" items="[[metadataTags]]">
+          <paper-item>[[item.tag]]</paper-item>
+        </template>
+        <paper-item>None</paper-item>
+      </paper-menu>
+    </paper-dropdown-menu>
+  </div>
+  <div class="control-holder">
     <div class="title">Upload</div>
     <paper-button raised class="text-button upload-button"
         on-click="_getFile">Choose File</paper-button>
@@ -194,20 +217,18 @@ svg.icon {
   </div>
   <div class="control-holder">
     <div class="title">Color</div>
-    <paper-dropdown-menu no-label-float no-animations noink class="color-dropdown">
-      <paper-menu class="dropdown-content" selected="{{_colorByIndex}}">
-        <paper-item>Structure</paper-item>
-        <paper-item>Device</paper-item>
-        <template is="dom-if" if="[[hasStats]]">
-          <paper-item>Compute time</paper-item>
-          <paper-item>Memory</paper-item>
-        </template>
-      </paper-menu>
-    </paper-dropdown-menu>
+    <paper-radio-group selected="{{colorBy}}">
+      <paper-radio-button name="structure">Structure</paper-radio-button>
+      <paper-radio-button name="device">Device</paper-radio-button>
+      <template is="dom-if" if="[[_statsNotNull(stats)]]">
+        <paper-radio-button name="compute_time">Compute time</paper-radio-button>
+        <paper-radio-button name="memory">Memory</paper-radio-button>
+      </template>
+    </paper-radio-group>
   </div>
   <div>
     <template is="dom-if" if="[[_isGradientColoring(colorBy)]]">
-      <svg width="160" height="20" style="margin: 0 5px" class="color-text">
+      <svg width="140" height="20" style="margin: 0 5px" class="color-text">
         <defs>
           <linearGradient id="linearGradient" x1="0%" y1="0%" x2="100%" y2="0%">
             <stop class="start" offset="0%"
@@ -216,7 +237,7 @@ svg.icon {
                 stop-color$="[[_currentGradientParams.endColor]]"/>
           </linearGradient>
         </defs>
-        <rect x="0" y="0" width="160" height="20" fill="url(#linearGradient)"
+        <rect x="0" y="0" width="135" height="20" fill="url(#linearGradient)"
             stroke="black" />
       </svg>
       <div class="domainValues color-text">
@@ -359,19 +380,22 @@ Polymer({
   is: 'tf-graph-controls',
   properties: {
     // Public API.
-    hasStats: {
-      type: Boolean
-    },
+    stats: Object,
     colorBy: {
       type: String,
+      value: 'structure',
       notify: true,
-      computed: '_getColorBy(_colorByIndex)'
+      readonly: true
     },
     colorByParams: Object,
     datasets: {
       type: Array,
       observer: '_datasetsChanged'
     },
+    metadataTags: {
+      type: Array,
+      computed: '_getMetadataTags(selectedDataset, datasets)'
+    },
     selectedDataset: {
       type: Number,
       notify: true,
@@ -382,18 +406,21 @@ Polymer({
       type: Object,
       notify: true
     },
-    // Private API.
-    _colorByIndex: {
+    selectedMetadataTag: {
       type: Number,
-      value: 0 // Defaults to 'structure'.
+      notify: true,
+      value: -1
     },
     _currentGradientParams: {
       type: Object,
       computed: '_getCurrentGradientParams(colorByParams, colorBy)'
     }
   },
-  _getColorBy: function(colorByIndex) {
-    return ["structure", "device", "compute_time", "memory"][colorByIndex];
+  _statsNotNull: function(stats) {
+    return stats != null;
+  },
+  _numSessionRuns: function(metadataTags) {
+    return metadataTags != null ? metadataTags.length : 0;
   },
   _getBackgroundColor: function(color) {
     return 'background-color:' + color;
@@ -446,8 +473,13 @@ Polymer({
       this._setDownloadFilename(this.datasets[this.selectedDataset].path);
     }
   },
+  _getMetadataTags: function(selectedDataset, datasets) {
+    return this.datasets[selectedDataset].runMetadata;
+  },
   _selectedDatasetChanged: function(newDataset, oldDataset) {
     if (this.datasets) {
+      this.set('selectedMetadataTag', -1);
+      this.set('colorBy', 'structure');
       this._setDownloadFilename(this.datasets[newDataset].path);
     }
   },
diff --git a/tensorflow/tensorboard/components/tf-graph/tf-graph-scene.html b/tensorflow/tensorboard/components/tf-graph/tf-graph-scene.html
index f01f5b752b..2e977d2699 100644
--- a/tensorflow/tensorboard/components/tf-graph/tf-graph-scene.html
+++ b/tensorflow/tensorboard/components/tf-graph/tf-graph-scene.html
@@ -676,12 +676,14 @@ Polymer({
     * UI controls.
     */
   _colorByChanged: function() {
-    // We iterate through each svg node and update its state.
-    _.each(this._nodeGroupIndex, function(nodeGroup, nodeName) {
-      this._updateNodeState(nodeName);
-    }, this);
-    // Notify also the minimap.
-    this.minimap.update();
+    if (this.renderHierarchy != null) {
+      // We iterate through each svg node and update its state.
+      _.each(this._nodeGroupIndex, function(nodeGroup, nodeName) {
+        this._updateNodeState(nodeName);
+      }, this);
+      // Notify also the minimap.
+      this.minimap.update();
+    }
   },
   fit: function() {
     tf.graph.scene.fit(this.$.svg, this.$.root, this._zoom, function() {
diff --git a/tensorflow/tensorboard/components/tf-graph/tf-graph.html b/tensorflow/tensorboard/components/tf-graph/tf-graph.html
index 7f29049ea6..d0bc9c55fb 100644
--- a/tensorflow/tensorboard/components/tf-graph/tf-graph.html
+++ b/tensorflow/tensorboard/components/tf-graph/tf-graph.html
@@ -43,7 +43,6 @@ paper-button {
           highlighted-node="[[_getVisible(highlightedNode)]]"
           selected-node="[[selectedNode]]"
           color-by="[[colorBy]]"
-          name="[[graphName]]"
           progress="[[progress]]"
     ></tf-graph-scene>
   </div>
@@ -63,6 +62,10 @@ Polymer({
       observer: '_graphChanged'
     },
     basicGraph: Object,
+    stats: {
+      type: Object,
+      observer: '_statsChanged'
+    },
     hierarchyParams: Object,
     progress: {
       type: Object,
@@ -101,6 +104,14 @@ Polymer({
   observers: [
     '_buildRenderHierarchy(graphHierarchy)'
   ],
+  _statsChanged: function(stats) {
+    if (stats != null) {
+      tf.graph.joinStatsInfoWithGraph(this.basicGraph, stats);
+      tf.graph.hierarchy.joinAndAggregateStats(this.graphHierarchy, stats);
+      // Recompute the rendering information.
+      this._buildRenderHierarchy(this.graphHierarchy);
+    }
+  },
   _buildRenderHierarchy: function(graphHierarchy) {
     tf.time('new tf.graph.render.Hierarchy', function() {
       if (graphHierarchy.root.type !== tf.graph.NodeType.META) {