[pt2] add metas for median ops#106272
[pt2] add metas for median ops#106272nkaretnikov wants to merge 4 commits intogh/nkaretnikov/161/basefrom
median ops#106272Conversation
[ghstack-poisoned]
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106272
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 3a4c698: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
From the list in #105105 |
| @register_meta(aten.nanmedian.default) | ||
| def meta_nanmedian(input): | ||
| @register_meta([aten.median.default, aten.nanmedian.default]) | ||
| def meta_median(input): |
There was a problem hiding this comment.
median is the same as nanmedian without NaN handling
[ghstack-poisoned]
| def meta_nanmedian_dim(input, dim=-1, keepdim=False): | ||
| def meta_median_dim(input, dim=-1, keepdim=False): | ||
| if device_hint(input) == "cuda": | ||
| utils.alert_not_deterministic("median CUDA with indices output") |
There was a problem hiding this comment.
Repro (this alert is only on CUDA):
PYTORCH_TEST_WITH_INDUCTOR=1 python -bb test/test_torch.py -v --use-pytest --import-slow-tests --import-disabled-tests -k test_nondeterministic_alert_median -v --capture=no
There was a problem hiding this comment.
This is kind of weird, we shouldn't be on the hook for reporting non-determinism in meta kernels...
|
The XLA failure is real (counters) and will require updating tests on the XLA side. |
[ghstack-poisoned]
|
Updated XLA pin to point to companion PR: pytorch/xla#5376 |
|
XLA people told me the issue on the XLA side needs to be fixed in a different way. Converting to a draft, so it's not merged by accident. UPD: The XLA team created a new PR that switches to a different function to test the fallback. Waiting until that one is merged. pytorch/xla#5393 UPD2: Updated XLA hash. |
[ghstack-poisoned]
Pull Request resolved: #106273 Approved by: https://github.com/ezyang ghstack dependencies: #106272
Stack from ghstack:
modeops #106273medianops #106272