From ef4e6036bcd89fbd1565e86abeffe36f58fd53d6 Mon Sep 17 00:00:00 2001
From: David Reiss <dreiss@fb.com>
Date: Wed, 19 May 2021 12:21:43 -0700
Subject: [PATCH] model_dump: Handle dict rendering (#57657)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/57657

Test Plan: Clicked around a model with some dicts in it.

Reviewed By: malfet

Differential Revision: D28531397

Pulled By: dreiss

fbshipit-source-id: 069690f147e91eadd76fec5f5ca4eec057abcb98
---
 torch/utils/model_dump/code.js | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/torch/utils/model_dump/code.js b/torch/utils/model_dump/code.js
index 4380ead425c..d5edb802ac7 100644
--- a/torch/utils/model_dump/code.js
+++ b/torch/utils/model_dump/code.js
@@ -124,6 +124,10 @@ class ModelData extends Component {
       // TODO: Maybe show simple lists and tuples on one line.
       return true;
     }
+    if (data.__is_dict__) {
+      // TODO: Maybe show simple (empty?) dicts on one line.
+      return true;
+    }
     if (data.__module_type__) {
       return true;
     }
@@ -133,7 +137,7 @@ class ModelData extends Component {
     if (data.__qtensor__) {
       return false;
     }
-    throw new Error("TODO: handle dict, etc.");
+    throw new Error("Can't handle data type.", data);
   }
 
   renderHeadline(data) {
@@ -159,6 +163,9 @@ class ModelData extends Component {
     if (data.__tuple_values__) {
       return "tuple((";
     }
+    if (data.__is_dict__) {
+      return "dict({";
+    }
     if (data.__module_type__) {
       return data.__module_type__ + "()";
     }
@@ -181,7 +188,7 @@ class ModelData extends Component {
       return this.renderTensor(
         "qtensor", dtype, key, device, numel, offset, size, stride, grad, extra_parts);
     }
-    throw new Error("TODO: handle dict, etc.");
+    throw new Error("Can't handle data type.", data);
   }
 
   renderTensor(
@@ -237,6 +244,18 @@ class ModelData extends Component {
       // Handled the same as lists.
       return this.renderBody(indent, data.__tuple_values__);
     }
+    if (data.__is_dict__) {
+      let new_indent = indent + "\u00A0\u00A0";
+      let parts = [];
+      for (let idx = 0; idx < data.keys.length; idx++) {
+        if (typeof(data.keys[idx]) != "string") {
+          parts.push(html`<br/>${new_indent}Non-string key`);
+        } else {
+          parts.push(html`<br/><${ModelData} prefix=${data.keys[idx] + ": "} indent=${new_indent} data=${data.values[idx]} />`);
+        }
+      }
+      return parts;
+    }
     if (data.__module_type__) {
       const mstate = data.state;
       if (mstate === null || typeof(mstate) != "object") {
@@ -245,6 +264,7 @@ class ModelData extends Component {
       let new_indent = indent + "\u00A0\u00A0";
       let parts = [];
       if (mstate.__is_dict__) {
+        // TODO: Less copy/paste between this and normal dicts.
         for (let idx = 0; idx < mstate.keys.length; idx++) {
           if (typeof(mstate.keys[idx]) != "string") {
             parts.push(html`<br/>${new_indent}Non-string key`);
@@ -267,7 +287,7 @@ class ModelData extends Component {
     if (data.__qtensor__) {
       throw "Should not reach here."
     }
-    throw new Error("TODO: handle dict, etc.");
+    throw new Error("Can't handle data type.", data);
   }
 
   render({data, indent, prefix}, {shown}) {