[GitHub] [incubator-tvm] tkonolige commented on a change in pull request #6854: [RELAY,TOPI] Add scatter_nd op

[GitBox <gi...@apache.org>]

tkonolige commented on a change in pull request #6854:
URL: https://github.com/apache/incubator-tvm/pull/6854#discussion_r518246012



##########
File path: python/tvm/relay/op/_tensor_grad.py
##########
@@ -803,3 +805,15 @@ def arange_grad(orig, grad):
     grad_step = cast_like(_sum(grad_step), step)
 
     return [grad_start, grad_stop, grad_step]
+
+
+@register_gradient("gather_nd")
+def gather_nd_grad(orig, grad):
+    data, indices = orig.args
+    return [scatter_nd(grad, indices, data.checked_type.concrete_shape), zeros_like(indices)]
+
+
+# @register_gradient("scatter_nd")

Review comment:
       I'm getting some really weird results when using gather_nd as the gradient. Values are like 1e303.




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org