[MPS] Add i0 op #137849
[MPS] Add i0 op #137849
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/137849
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit af0f1f2 with merge base bc232e3 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/mps
Attention! native_functions.yaml was changedIf you are adding a new function or defaulted argument to native_functions.yaml, you cannot use it from pre-existing Python frontend code until our FC window passes (two weeks). Split your PR into two PRs, one which adds the new C++ functionality, and one that makes use of it from Python, and land them two weeks apart. See https://github.com/pytorch/pytorch/wiki/PyTorch's-Python-Frontend-Backward-and-Forward-Compatibility-Policy#forwards-compatibility-fc for more info. Caused by: |
| 6.88975834691682398426E-5, 3.36911647825569408990E-3, | ||
| 8.04490411014108831608E-1}; | ||
|
|
||
| return static_cast<T>((exp(x) * chbevl(32.0 / x - 2.0, B, 25)) / sqrt(x)); |
There was a problem hiding this comment.
| return static_cast<T>((exp(x) * chbevl(32.0 / x - 2.0, B, 25)) / sqrt(x)); | |
| return static_cast<T>((exp(x) * chbevl(32.0 / x - 2.0, B, 25)) * rsqrt(x)); |
Not sure if this affects the numeric, but it should have a special rsqrt instruction we can exploit
There was a problem hiding this comment.
Compiler should be able to figure out that 1.0 / sqrt could be replaced with rsqrt, but sure multiplication of 3 elements indeed sounds nice
There was a problem hiding this comment.
Actually shouldn't be able to since the numerics differ ever so slightly apparently.
To match behavior for torch.special.i0 Noticed while looking at the failures in #137849
|
@pytorchbot merge -f "Lint + MPS are green" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
To match behavior of `torch.special.i0` Noticed while looking at the failures in #137849 Also, add explicit high-precision template specialization for `calc_i0` and `calc_i1` for `BFloat16` and `Half` Pull Request resolved: #137899 Approved by: https://github.com/Skylion007
3 participants
More-or-less verbatim copy of
pytorch/aten/src/ATen/native/Math.h
Line 101 in 47c8aa8
Plus a bit of a MPS boilerplate code
Update test_mps.py to mark kaiser_window and i0 as passing