From 83e4e2dbf5e6fd2e8ecb02ebc6817b8ab3d20dee Mon Sep 17 00:00:00 2001
From: Qiumin Xu <qiuminxu@google.com>
Date: Wed, 11 Jul 2018 14:36:54 -0700
Subject: [PATCH 1/2] Fix memory viewer due to minor compiler changes.

---
 .../buffer_details/buffer-details.html        |  4 +--
 .../profile/memory_viewer/utils/utils.ts      |  2 ++
 .../profile/memory_viewer/xla/shape.ts        | 25 +++++++++++--------
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/tensorboard/plugins/profile/memory_viewer/buffer_details/buffer-details.html b/tensorboard/plugins/profile/memory_viewer/buffer_details/buffer-details.html
index b782457164c..6648afc5fcf 100644
--- a/tensorboard/plugins/profile/memory_viewer/buffer_details/buffer-details.html
+++ b/tensorboard/plugins/profile/memory_viewer/buffer_details/buffer-details.html
@@ -19,7 +19,7 @@
 <!--
   tf-mv-bar shows a small % utilization bar, colored using flameColor.
 -->
-<dom-module id="tf-nv-bar">
+<dom-module id="tf-mv-bar">
   <style>
 :host {
   display: inline-block;
@@ -74,7 +74,7 @@ <h4>Size: <span>[[size]]</span><b> MiB</b></h4>
           <tf-mv-bar value="[[utilization]]"></tf-mv-bar>
         </div>
         <div hidden="[[!node.shape]]">
-          <b>Shape: </b>
+          <b>Shape (and minor-to-major order): </b>
           <code class="expression">[[node.shape]]</code>
         </div>
         <div hidden="[[!node.tfOpName]]">
diff --git a/tensorboard/plugins/profile/memory_viewer/utils/utils.ts b/tensorboard/plugins/profile/memory_viewer/utils/utils.ts
index 02809602af0..0d59c19c97e 100644
--- a/tensorboard/plugins/profile/memory_viewer/utils/utils.ts
+++ b/tensorboard/plugins/profile/memory_viewer/utils/utils.ts
@@ -51,6 +51,8 @@ export function byteSizeOfPrimitiveType(type: string): number {
       return 8;
     case 'C64':
       return 8;
+    case 'TOKEN':
+      return 0;
     default:
       console.error('Unhandled primitive type ' + type);
       return 0;
diff --git a/tensorboard/plugins/profile/memory_viewer/xla/shape.ts b/tensorboard/plugins/profile/memory_viewer/xla/shape.ts
index 602ba287f2f..15ef127de75 100644
--- a/tensorboard/plugins/profile/memory_viewer/xla/shape.ts
+++ b/tensorboard/plugins/profile/memory_viewer/xla/shape.ts
@@ -48,17 +48,20 @@ export class Shape {
 
     // We make a simplifying assumption here that the minimum size of a tuple
     // member is int64.
+    if (this.elementType === 'TOKEN') {
+      return 0;
+    }
     if (this.elementType === 'TUPLE') {
       return INT64_BYTES * this.tupleShapes.length;
     }
     let byteSize = 0;
-    if (this.layout.format == 'DENSE') {
+    // We assume the layout format is 'DENSE' by default.
+    if (!this.layout || this.layout.format == 'DENSE') {
       const allocatedElementCount =
           this.dimensions.reduce((count, item) => count * item, 1);
       byteSize += allocatedElementCount *
           memory_viewer.byteSizeOfPrimitiveType(this.elementType);
-    }
-    if (this.layout.format == 'SPARSE') {
+    }else if (this.layout.format == 'SPARSE') {
       const maxElements = parseInt(this.layout.maxSparseElements, 10);
       byteSize = maxElements * memory_viewer.byteSizeOfPrimitiveType(this.elementType);
 
@@ -82,16 +85,16 @@ export class Shape {
       }
       text += ')';
       return text;
-    } else {
-      let result = this.elementType.toLowerCase() + '[';
-      result += this.dimensions.join() + ']';
-      if (!(this.elementType === 'OPAQUE') && this.dimensions.length > 0) {
-        if (this.layout) {
-          result += this.humanLayoutString(this.layout);
-        }
+    }
+    let result = this.elementType.toLowerCase() + '[';
+    result += this.dimensions.join() + ']';
+    if (!(this.elementType === 'OPAQUE') && !(this.elementType === 'TOKEN') &&
+        this.dimensions.length > 0) {
+      if (this.layout) {
+        result += this.humanLayoutString(this.layout);
       }
-      return result;
     }
+    return result;
   }
 
   /**

From 65a4fd9d76449a26ce8def861db3a5530a401a1c Mon Sep 17 00:00:00 2001
From: Qiumin Xu <qiuminxu@google.com>
Date: Wed, 11 Jul 2018 17:50:15 -0700
Subject: [PATCH 2/2] fix comments

---
 tensorboard/plugins/profile/memory_viewer/xla/shape.ts | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorboard/plugins/profile/memory_viewer/xla/shape.ts b/tensorboard/plugins/profile/memory_viewer/xla/shape.ts
index 15ef127de75..df1075512cb 100644
--- a/tensorboard/plugins/profile/memory_viewer/xla/shape.ts
+++ b/tensorboard/plugins/profile/memory_viewer/xla/shape.ts
@@ -46,11 +46,11 @@ export class Shape {
   unpaddedHeapSizeBytes(): number {
     const INT64_BYTES = 8;
 
-    // We make a simplifying assumption here that the minimum size of a tuple
-    // member is int64.
     if (this.elementType === 'TOKEN') {
       return 0;
     }
+    // We make a simplifying assumption here that the minimum size of a tuple
+    // member is int64.
     if (this.elementType === 'TUPLE') {
       return INT64_BYTES * this.tupleShapes.length;
     }
@@ -63,7 +63,8 @@ export class Shape {
           memory_viewer.byteSizeOfPrimitiveType(this.elementType);
     }else if (this.layout.format == 'SPARSE') {
       const maxElements = parseInt(this.layout.maxSparseElements, 10);
-      byteSize = maxElements * memory_viewer.byteSizeOfPrimitiveType(this.elementType);
+      byteSize = maxElements *
+          memory_viewer.byteSizeOfPrimitiveType(this.elementType);
 
       // Add byte size of sparse indices, assume each indice is int64 type.
       byteSize += maxElements * this.dimensions.length * INT64_BYTES;