Print requires_grad and grad_fn in string repr of tensor (pytorch#8211)

colesbury · weiyangfb · commit 0bf9a0dba170 · 2018-06-11T12:17:26.000-07:00
For example:

  &gt;&gt;&gt; torch.ones(3).requires_grad_()
  tensor([ 1.,  1.,  1.], requires_grad=True)

  &gt;&gt;&gt; torch.ones(3).requires_grad_() * 5
  tensor([ 5.,  5.,  5.], grad_fn=&lt;MulBackward0&gt;)

The suffix (dtype, requires_grad, grad_fn) wraps to a new line if
it would cause the the line to exceed the linewidth.

  &gt;&gt;&gt; torch.ones(10).double().requires_grad_()
  tensor([ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.],
         dtype=torch.float64, requires_grad=True)
diff --git a/torch/_tensor_str.py b/torch/_tensor_str.py
@@ -196,6 +196,14 @@ def _tensor_str(self, indent, fmt, scale, sz, summarize):
     return '[' + tensor_str + ']'
 
 
+def _maybe_wrap_suffix(suffix, indent, tensor_str):
+    suffix_len = len(suffix)
+    last_line_len = len(tensor_str) - tensor_str.rfind('\n') + 1
+    if suffix_len > 2 and last_line_len + suffix_len > PRINT_OPTS.linewidth:
+        return ',\n' + ' ' * indent + suffix[2:]
+    return suffix
+
+
 def get_summarized_data(self):
     dim = self.dim()
     if dim == 0:
@@ -224,27 +232,36 @@ def _str(self):
     indent = len(prefix)
     summarize = self.numel() > PRINT_OPTS.threshold
 
-    suffix = ')'
+    suffix = ''
     if not torch._C._is_default_type_cuda():
         if self.device.type == 'cuda':
-            suffix = ', device=\'' + str(self.device) + '\'' + suffix
+            suffix += ', device=\'' + str(self.device) + '\''
     else:
         if self.device.type == 'cpu' or torch.cuda.current_device() != self.device.index:
-            suffix = ', device=\'' + str(self.device) + '\'' + suffix
+            suffix += ', device=\'' + str(self.device) + '\''
 
     if self.numel() == 0:
         # In an empty tensor, there are no elements to infer if the dtype should be int64,
         # so it must be shown explicitly.
         if self.dtype != torch.get_default_dtype():
-            suffix = ', dtype=' + str(self.dtype) + suffix
+            suffix += ', dtype=' + str(self.dtype)
         tensor_str = '[]'
     else:
         if self.dtype != torch.get_default_dtype() and self.dtype != torch.int64:
-            suffix = ', dtype=' + str(self.dtype) + suffix
+            suffix += ', dtype=' + str(self.dtype)
 
         fmt, scale, sz = _number_format(get_summarized_data(self) if summarize else self)
         if scale != 1:
             prefix = prefix + SCALE_FORMAT.format(scale) + ' ' * indent
         tensor_str = _tensor_str(self, indent, fmt, scale, sz, summarize)
 
+    if self.grad_fn is not None:
+        suffix += ', grad_fn=<{}>'.format(type(self.grad_fn).__name__)
+    elif self.requires_grad:
+        suffix += ', requires_grad=True'
+
+    suffix += ')'
+
+    suffix = _maybe_wrap_suffix(suffix, indent, tensor_str)
+
     return prefix + tensor_str + suffix
