viable/strict/1791919581: [MPS] Fix precision loss in complex sqrt (#200736)
- PyTorch: 2284 events in the last 90 days
- PyTorch: 2254th Release in the last 90 days
- Previous: earlier the same day · viable/strict/1791916658: [Stable-ABI] Add torch::stable::masked_select. (#198048)
What happened
The Metal complex sqrt computed the imaginary part as sqrt((|z| - re) / 2) , which suffers catastrophic cancellation when |im| (and the real part likewise via |z| + re when re ). E.g. on MPS sqrt(1+1e-4j) returned 1+0j instead of 1+5e-5j for complex64, and chalf sqrt(0.7783+0.0038j) had a 3.4x error in the imaginary part, which made test_sqrt fail once its explicit seed was removed. Fix by adding a c10::metal::sqrt…
Summary assembled by rule from the sources below