Dispatch to mv rather than mm in the case that tensor1.ndim == 1 and tensor2.ndim == 2
This should hopefully be faster, it makes the calling code simpler, and
it solves a bug when using matmul with the out= parameter (before it
would throw an incorrect error).
Pull Request resolved: https://github.com/pytorch/pytorch/pull/75195
Approved by: https://github.com/ezyang
diff --git a/aten/src/ATen/native/LinearAlgebra.cpp b/aten/src/ATen/native/LinearAlgebra.cpp
index 1fc6e83..1f96e78 100644
--- a/aten/src/ATen/native/LinearAlgebra.cpp
+++ b/aten/src/ATen/native/LinearAlgebra.cpp
@@ -1661,8 +1661,7 @@
} else if (dim_tensor1 == 2 && dim_tensor2 == 1) {
return has_out ? at::mv_out(out, tensor1, tensor2) : tensor1.mv(tensor2);
} else if (dim_tensor1 == 1 && dim_tensor2 == 2) {
- return has_out ? at::mm_out(out, tensor1.unsqueeze(0), tensor2).squeeze_(0)
- : tensor1.unsqueeze(0).mm(tensor2).squeeze_(0);
+ return has_out ? at::mv_out(out, tensor2.t(), tensor1) : tensor2.t().mv(tensor1);
} else if (dim_tensor1 == 2 && dim_tensor2 == 2) {
return has_out ? at::mm_out(out, tensor1, tensor2) : tensor1.mm(tensor2);
} else if (should_fold_into_mm(tensor1, tensor2)) {
