🐛 Describe the bug

When using the MPS device on an Apple M2 Pro, I see the following behavior:

import torch

t = torch.tensor([1,2,3]).to("mps")
t.clamp_min(1.0)

Output: tensor([0., 0., 0.], device='mps:0')

If the correct type is used for the clamp value:

t.clamp_min(1)

Output: tensor([1, 2, 3], device='mps:0')

Also note, the same behavior is seen for the clamp method:

t.clamp(min=1.)

Output: tensor([0., 0., 0.], device='mps:0')
with the following message:

miniconda3/envs/test2/lib/python3.12/site-packages/torch/_tensor_str.py:138: UserWarning: MPS: nonzero op is supported natively starting from macOS 13.0. Falling back on CPU. This may have performance implications. (Triggered internally at /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/mps/operations/Indexing.mm:335.)
  nonzero_finite_vals = torch.masked_select(

This seems to only occur on the MPS device, whereas CPU (on same machine) and cuda (different environment) produced the correct output.

Versions

Collecting environment information...
PyTorch version: 2.4.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.6.7 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.1.0.2.5)
CMake version: version 3.27.6
Libc version: N/A

Python version: 3.12.4 | packaged by Anaconda, Inc. | (main, Jun 18 2024, 10:07:17) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-13.6.7-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Pro

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.4.0
[conda] numpy 1.26.4 pypi_0 pypi
[conda] torch 2.4.0 pypi_0 pypi

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen