Dispatch numpy.take_along_axis to torch.take_along_dim
#108880
Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/108880
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit 1173859 with merge base bde75eb ( UNSTABLE - The following jobs failed but were likely due to flakiness present on trunk and has been marked as unstable:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
test/dynamo/test_misc.py
Outdated
| torch.float32, | ||
| requires_grad=False, | ||
| include_0d=False, | ||
| include_empty=False, |
There was a problem hiding this comment.
What's the issue with the empty one?
There was a problem hiding this comment.
It fails in a test case on test_misc.py as the indices has dtype uint8_t:
...
torch/_dynamo/utils.py:1339: in <lambda>
lambda: run_node(tx.output, node, args, kwargs, nnmodule)
torch/_dynamo/utils.py:1411: in run_node
raise RuntimeError(fn_str + str(e)).with_traceback(e.__traceback__) from e
torch/_dynamo/utils.py:1398: in run_node
return node.target(*args, **kwargs)
torch/_dynamo/utils.py:1699: in __call__
out = self.f(*args, **kwargs)
torch/_numpy/_normalizations.py:213: in wrapped
result = func(*args, **kwds)
torch/_numpy/_funcs_impl.py:889: in take_along_axis
return torch.take_along_dim(arr, indices, axis)
E torch._dynamo.exc.TorchRuntimeError: Failed running call_function <Wrapped function <original take_along_axis>>(*(FakeTensor(..., size=(s0,)), FakeTensor(..., size=(0,), dtype=torch.uint8), 0), **{}):
E torch.take_along_dim(): dtype of indices should be Long but got Byte
E
E from user code:
E File "/home/guilhermeleobas/git/pytorch/test/dynamo/test_misc.py", line 1283, in fn
E return np.take_along_axis(x, i, a)
E
E Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
E
E
E You can suppress this exception and fall back to eager by setting:
E import torch._dynamo
E torch._dynamo.config.suppress_errors = True
E
E
E To execute this test, run the following from the base repo dir:
E python test/dynamo/test_misc.py -k test_numpy_take_along_axis
E
E This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0
There was a problem hiding this comment.
Makes sense. Actually, it may be best to use sample_inputs_take_along_dim. Note that the order of the args for indices and dims is swapped wrt. gather.
There was a problem hiding this comment.
Are you ok if I skip the tests from sample_inputs_take_along_dim that don't include dim?
There was a problem hiding this comment.
What do you mean exactly?
There was a problem hiding this comment.
Sorry for not including more information in my previous comment. There are two tests that don't include the dim argument in sample_inputs_take_along_dim
There was a problem hiding this comment.
Actually, Numpy also has this behaviour, so we should also implement it. Can you modify the decomposition accordingly?
Also, could you send a separate PR fixing the docs of this operation noting that dim can also be None?
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
There was a problem hiding this comment.
Just a small note, feel free to merge once you change the sampling function.
You can even clean-up the include_0d and the include_empty args for sample_inputs_gather as these will not be used anywhere any more.
Feel free to merge once you change this. In a stack, if you call pytorchbot merge on a given PR, it'll merge that PR and all the previous ones.
test/dynamo/test_misc.py
Outdated
| torch.float32, | ||
| requires_grad=False, | ||
| include_0d=False, | ||
| include_empty=False, |
There was a problem hiding this comment.
Makes sense. Actually, it may be best to use sample_inputs_take_along_dim. Note that the order of the args for indices and dims is swapped wrt. gather.
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
|
Also, note that you can even clean-up the include_0d and the include_empty args for sample_inputs_gather as these will not be used anywhere any more. |
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng Xia-Weiwen wenzhe-nrv jiayisunx chenyang78 aakhundov kadeng [ghstack-poisoned]
Pull Request resolved: #109120 Approved by: https://github.com/lezcano ghstack dependencies: #108879, #108880
Pull Request resolved: #109125 Approved by: https://github.com/lezcano ghstack dependencies: #108879, #108880, #109120
4 participants
Stack from ghstack (oldest at bottom):
dim=Nonecase #109120numpy.take_along_axistotorch.take_along_dim#108880SymIntsupport totorch.take_along_dim#108879cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng