t-vi commented on a change in pull request #6740:
URL: https://github.com/apache/incubator-tvm/pull/6740#discussion_r510653148
##########
File path: python/tvm/relay/frontend/pytorch.py
##########
@@ -2357,6 +2357,36 @@ def _impl(inputs, input_types):
return _impl
+def _bincount():
+ def _impl(inputs, input_types):
+ data = inputs[0]
+ weights = inputs[1]
+ maximum = _op.max(data)
+ dim = maximum + _expr.const(1, dtype="int64")
+ if weights:
+ out_dtype = "float32"
Review comment:
While int64 is fixed in PyTorch without weights, passing in float64
weights returns float64, so we should probably take the weight's dtype.
##########
File path: tests/python/frontend/pytorch/test_forward.py
##########
@@ -3147,18 +3147,28 @@ def __init__(self, dim=0):
def forward(self, data, index, src):
return torch.scatter(data, dim=self.dim, index=index, src=src)
+ class ScatterAdd(Module):
+ def __init__(self, dim=0):
+ super().__init__()
+ self.dim = dim
+
+ def forward(self, data, index, src):
+ return torch.scatter_add(data, dim=self.dim, index=index, src=src)
+
Review comment:
As a general remark (rather than necessarily to act upon), I had hoped
that supporting traced functions would help us make the test suite more lean...
Somehow the module creation for each function looks unwieldy to me.
----------------------------------------------------------------
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:
[email protected]