commit 287f4349008536dde3c276f1c25024fb129a9d02
author Junjie Bai <bai@in.tum.de> Fri Sep 22 19:51:39 2017 -0700
committer Edward Z. Yang <ezyang@mit.edu> Sat Sep 23 11:17:27 2017 -0400
tree 0bfde6966877a00f14de7e9e9e3466a342820661
parent 767f704b84b324f89bc7f3f408aab23f3e84f835

Add support for exporting Addmm with alpha != 1 or beta != 1

torch/autograd/_functions/blas.py

diff --git a/torch/autograd/_functions/blas.py b/torch/autograd/_functions/blas.py
index c7bcb3b2..787ac3b 100644
--- a/torch/autograd/_functions/blas.py
+++ b/torch/autograd/_functions/blas.py
@@ -17,10 +17,10 @@
 
     @staticmethod
     def symbolic(g, add_matrix, matrix1, matrix2, alpha=1, beta=1, inplace=False):
-        # TODO: manually insert the necessary scaling, since ONNX doesn't
-        # natively support it
-        if alpha != 1 or beta != 1:
-            return None
+        if alpha != 1:
+            matrix1 = g.op("Scale", matrix1, scale_f=alpha)
+        if beta != 1:
+            add_matrix = g.op("Scale", add_matrix, scale_f=beta)
         # TODO: Talk to ONNX about why their FC involves a transpose
         matrix2_t = g.op("Transpose", matrix2)
         return g.op("FC", matrix1, matrix2_t, add_matrix)