pytorch

f7e7a15a - Fix `NaN` comparison in `torch.median` (#38216)

Fix `NaN` comparison in `torch.median` (#38216) Summary: Fixes https://github.com/pytorch/pytorch/issues/38018 when calling `eq_with_nan(v, kValue)` having `v` and `kValue` both `nan` is returning `false` when it should be `true`. https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/SortingKthValue.cu#L76 The implementation is using intrinsics such as `__double_as_longlong` and comparing their bit representations. But the values of the bits obtained for both nans are different. `9221120237041090560` for `v` `9223372036854775807` for `kValue` two different nans have different bit representations, so we have to do additional comparisons to fix this. I changed this comparison and it seems to be working now. However, when compared to a CPU implementation, the returned indices for the values seems to be random but valid. Probably this is an effect of the comparison order in the Cuda version. I am not sure if this is ok since all the indices point to valid elements. For the snippet in the issue I get the following: ``` # CUDA Values tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan], device='cuda:0', dtype=torch.float64) # CUDA indices tensor([304, 400, 400, 528, 304, 304, 528, 336, 304, 432, 400, 280, 280, 336, 304, 336, 400, 304, 336, 560], device='cuda:0') ``` ``` # CPU values tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan], dtype=torch.float64) # CPU indices tensor([515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515, 515]) ``` Also, maybe its better to change the `eq_with_nan` implementations to address this instead? I am not sure if this will cause code to break in other places though ... Pull Request resolved: https://github.com/pytorch/pytorch/pull/38216 Differential Revision: D21517617 Pulled By: ngimel fbshipit-source-id: deeb7bb0ac519a03aa0c5f365005a9150e6404e6