fix small bug about dense_grad

apache · May 29, 2020 · 331bea8 · 331bea8
1 parent 95b3ad9
commit 331bea8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/python/tvm/relay/op/_tensor_grad.py b/python/tvm/relay/op/_tensor_grad.py
@@ -472,8 +472,8 @@ def bias_add_grad(orig, grad):
 def dense_grad(orig, grad):
     """Returns [grad' @ weight, data @ grad']"""
     data, weight = orig.args
-    return [collapse_sum_like(transpose(grad) * weight, data),
-            collapse_sum_like(data * transpose(grad), weight)]
+    return [collapse_sum_like(_nn.dense(grad, transpose(weight)), data),
+            collapse_sum_like(_nn.dense(transpose(grad), transpose(data)), weight)]
 
 
 @register_gradient("reshape")

diff --git a/tests/python/relay/test_op_grad_level2.py b/tests/python/relay/test_op_grad_level2.py
@@ -162,6 +162,7 @@ def verify_dense_grad(d_shape, w_shape):
 def test_dense_grad():
     verify_dense_grad((1, 8), (16, 8))
     verify_dense_grad((1, 4), (3, 4))
+    verify_dense_grad((5, 4), (3, 4))
 
 
 def verify_batch_flatten_grad(d_shape):