Addressed comments

zou3519 · zou3519 · commit 35f90b2cbcd2 · 2018-02-01T16:04:32.000-05:00
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -2394,7 +2394,11 @@ class dont_convert(tuple):
     ('dist', (), ((S, S, S), 4), 'scalar_4_broadcast_lhs'),
     ('diag', (M, M), NO_ARGS, '2d'),
     ('diag', (3, 5), NO_ARGS, '2d_wide'),
+    ('diag', (3, 5), (2,), '2d_wide_pos'),
+    ('diag', (3, 5), (-2,), '2d_wide_neg'),
     ('diag', (5, 3), NO_ARGS, '2d_tall'),
+    ('diag', (5, 3), (2,), '2d_tall_pos'),
+    ('diag', (5, 3), (-2,), '2d_tall_neg'),
     ('diag', (M,), NO_ARGS, '1d'),
     ('diag', (M, M), (1,), '2d_1'),
     ('diag', (M, M), (2,), '2d_2'),
diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -192,7 +192,7 @@
   self: _det_with_svd_backward(grads, self, result0, result1, result2, result3)
 
 - name: diag(Tensor self, int64_t diagonal)
-  self: diag_backward(grad, self, diagonal)
+  self: diag_backward(grad, self.sizes(), diagonal)
 
 - name: dist(Tensor self, Tensor other, Scalar p)
   self: norm_backward(grad, self - other, p, result)
diff --git a/tools/autograd/templates/Functions.cpp b/tools/autograd/templates/Functions.cpp
@@ -711,25 +711,36 @@ Tensor smooth_l1_loss_double_backward_grad_output(const Tensor & grad, const Ten
   return (r * grad).sum().view({1});
 }
 
-Tensor diag_backward(const Tensor & grad, const Tensor & self, int64_t diagonal) {
-  auto ndimension = self.ndimension();
+static inline int64_t diag_size(int64_t height, int64_t width, int64_t diagonal) {
+  if (width > height) {
+    return diag_size(width, height, -diagonal);
+  }
+  // Assumes height >= width
+  auto longest_diag = width;
+  if (diagonal >= 0) {
+    return longest_diag - diagonal;
+  }
+  if (longest_diag < height + diagonal) {
+    return longest_diag;
+  }
+  return height + diagonal;
+}
+
+Tensor diag_backward(const Tensor & grad, IntList input_sizes, int64_t diagonal) {
+  auto ndimension = input_sizes.size();
   TORCH_ASSERT(ndimension == 1 || ndimension == 2);
 
-  auto grad_input = grad.diag(diagonal);
-  if (ndimension == 1 || self.size(0) == self.size(1)) {
-    return grad_input;
+  if (ndimension == 1 || input_sizes[0] == input_sizes[1]) {
+    return grad.diag(diagonal);
   }
 
-  // cat rows or cols to grad_input so that it matches self's shape.
-  auto length = grad_input.size(0);
-  auto self_nrows = self.size(0);
-  auto self_ncols = self.size(1);
-  if (self_nrows == length) {
-    auto extra_cols = grad_input.type().zeros({self_nrows, self_ncols - length});
-    return at::cat({grad_input, extra_cols}, 1);
-  }
-  auto extra_rows = grad_input.type().zeros({self_nrows - length, self_ncols});
-  return at::cat({grad_input, extra_rows});
+  // Input was a matrix but was not square
+  auto grad_input = grad.type().zeros(input_sizes);
+  auto diagonal_size = diag_size(input_sizes[0], input_sizes[1], diagonal);
+  auto storage_offset = diagonal >= 0 ? diagonal : -diagonal * input_sizes[1];
+  auto diag = grad_input.as_strided({diagonal_size}, {input_sizes[1] + 1}, storage_offset);
+  diag.copy_(grad);
+  return grad_input;
 }
 
 Tensor max_pool2d_double_backward(const Tensor & grad, const Tensor & indices) {
