fix small bug about dense_grad (apache#5695)

neo-ai · Jun 18, 2020 · 68c6501 · 68c6501
1 parent 5c916fc
commit 68c6501
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/python/tvm/relay/op/_tensor_grad.py b/python/tvm/relay/op/_tensor_grad.py
@@ -472,9 +472,10 @@ 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),
+                                        units=weight.checked_type.shape[1]), data),
+            collapse_sum_like(_nn.dense(transpose(grad), transpose(data),
+                                        units=data.checked_type.shape[1]), weight)]
 
 @register_gradient("reshape")
 def reshape_grad(orig, grad):

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):