Avoid NaN values in torch.cdist backward for p<1

[kurtamohler]

Fixes #36493

[codecov]

Codecov Report

Merging #45720 into master will decrease coverage by 0.07%.
The diff coverage is 61.70%.

@@            Coverage Diff             @@
##           master   #45720      +/-   ##
==========================================
- Coverage   68.61%   68.53%   -0.08%     
==========================================
  Files         405      409       +4     
  Lines       52045    52553     +508     
==========================================
+ Hits        35710    36018     +308     
- Misses      16335    16535     +200     

Impacted Files	Coverage Δ
torch/_torch_docs.py	100.00% <ø> (ø)
torch/cuda/amp/grad_scaler.py	23.94% <0.00%> (-0.46%)	⬇️
torch/nn/modules/loss.py	97.76% <ø> (ø)
torch/onnx/symbolic_opset9.py	35.56% <ø> (ø)
torch/overrides.py	97.08% <ø> (ø)
torch/quantization/fx/pattern_utils.py	87.23% <ø> (-1.86%)	⬇️
torch/onnx/symbolic_opset11.py	21.42% <9.09%> (-1.28%)	⬇️
.../testing/_internal/distributed/distributed_test.py	30.29% <12.19%> (-0.41%)	⬇️
torch/onnx/symbolic_opset10.py	40.00% <16.66%> (+3.74%)	⬆️
torch/utils/benchmark/utils/timer.py	78.82% <25.00%> (-8.85%)	⬇️
... and 26 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6acd7b6...476adff. Read the comment docs.

[albanD]

Looks mostly good.
Added comments inline.

[albanD]

Looks good.
We can think of more "near to zero" optimization in a follow up PR if we want to.

[facebook-github-bot]

@albanD has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

[nikitaved]

@alban, we can modify the formula for different cases p < 1, p > 1 to make it more numerically stable. It will be bc-breaking for sure.

[albanD]

Not all BC-breaking are a nogo. If we change the computed value to be closer to the correct thing, that's an OK thing to do.

The main thing is that we should make sure that it does not make other regions significantly worse.

[facebook-github-bot]

@albanD merged this pull request in 54aaffb.