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…
Background
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.