Don't end on inplace operators in einsum (#22111)

t-vi · facebook-github-bot · commit f5df0c9104b0 · 2019-06-24T20:39:20.000-07:00
Summary: Returning the result of an inplace `squeeze_` in `einsum` (which itself is traced) interacts badly with `autograd.Function`. I must admit that I'm not 100% certain whether it should be necessary to change this, but I consider this a good change overall. Fixes: #22072 Pull Request resolved: #22111 Differential Revision: D15974990 Pulled By: soumith fbshipit-source-id: 477e7f23833f02999085f665c175d062e7d32acd
diff --git a/aten/src/ATen/native/Linear.cpp b/aten/src/ATen/native/Linear.cpp
@@ -119,9 +119,13 @@ static Tensor sumproduct_pair(const Tensor& left_, const Tensor& right_, IntArra
 
   // finally squeeze summed dimensions if desired
   if (! keepdim) {
-    for (int i = dim-1; i>=0; i--)
-      if (sum_dims[i])
-        result.squeeze_(i);
+    auto sizes = result.sizes().vec();
+    for (int i = dim-1; i>=0; i--) {
+      if (sum_dims[i]) {
+        sizes.erase(sizes.begin() + i);
+      }
+    }
+    result = result.view(sizes);
   }
   return result;
 }
@@ -354,8 +358,12 @@ Tensor einsum(std::string eqn, TensorList tensors) {
     result = at::native::sumproduct_pair(result, preprocessed_operands[i], sum_dims, true);
   }
   // finally, we squeeze out all non-result dimensions
-  for (int64_t dim = num_total_idxes-1; dim >= num_output_dims; dim--)
-    result.squeeze_(dim);
+  auto sizes = result.sizes().vec();
+  for (int64_t dim = num_total_idxes-1; dim >= num_output_dims; dim--) {
+    sizes.erase(sizes.begin() + dim);
+  }
+
+  result = result.view(sizes);
   return result;
 }
 
