With #56872 landed, now functions registered in
native_functions.yaml will have automatically generated check about whether tensor parameters are on the same devices by default.
If your function allows tensors on different devices (e.g. using
TensorIterator ), you can set
device_check: NoCheck explicitly in
native_functions.yaml . See also https://github.com/…/tree/master/aten/src/ATen/native…
Many functions check whether the tensor parameters are on the same device with a few different implementations, which cause code duplications and confusions:
The Appendix section in https://github.com/pytorch/pytorch/pull/55241 summarizes a few existing device check implementation in codebase (there are more in codebase! ).
This is also raised by community in Cleaner mechanism in the source-code to check if multiple tensors are on the same device · Issue #56632 · pytorch/pytorch · GitHub.
This is a first step towards unifying device check code.