[MPS] Native nonzero implementation

[DenisVieriu97]

Fixes https://github.com/pytorch/pytorch/issues/124850

Replace previous MPSGraph nonzero construction with native nonzero op. For older OSes, fallback to CPU (previous implementation was not reliable and was comparable to CPU in speed).

cc @skotapati

[pytorch-bot[bot]]

:link: Helpful Links

:test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/125355

• :page_facing_up: Preview Python docs built from this PR
• :page_facing_up: Preview C++ docs built from this PR
• :question: Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

:white_check_mark: No Failures

As of commit 2f3d603493c784916b21b9bd65ab87b5d8d7d191 with merge base 1eb7b8eb6085322c96ba0bbacb039ef6f3b2875b (): :green_heart: Looks good so far! There are no failures yet. :green_heart:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[kulinseth]

=================================== FAILURES ===================================
_________ TestConsistencyCPU.test_output_match_argwhere_cpu_complex64 __________
Unexpected success
-----------------

@DenisVieriu97 , this test is passing on MacOS 13 as well.

diff --git a/test/test_mps.py b/test/test_mps.py
index 817a93fbb9e..34cc9663008 100644
--- a/test/test_mps.py
+++ b/test/test_mps.py
@@ -344,7 +344,6 @@ def mps_ops_modifier(ops):
         'any',
         'addcdiv',
         'addcmul',
-        'argwhere',
         'asin',
         'atan',
         'atanh',

Lets remove it from the After macOS 14 test list.

[DenisVieriu97]

@pytorchbot merge -g

[pytorch-bot[bot]]

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: unrecognized arguments: -g

usage: @pytorchbot [-h] {merge,revert,rebase,label,drci,cherry-pick,close} ...

Try @pytorchbot --help for more info.

[DenisVieriu97]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status here

[atalman]

@pytorchbot cherry-pick --onto release/2.3 -c critical

[pytorchbot]

Cherry picking #125355

The cherry pick PR is at https://github.com/pytorch/pytorch/pull/126100 and it is recommended to link a critical cherry pick PR with an issue

Details for Dev Infra team

Raised by workflow job