Bug when indexing 2D tensors using an MPS device

[joaquimgomez]

🐛 Describe the bug

When executing the following code using an MPS device:

import torch

device = torch.device("mps")

x = torch.tensor([[1, 0, 2, 0, -100, -100, -100],
                  [1, 0, 2, 0, -100, -100, -100]])
mask = torch.tensor([[True, True, True, True, False, False, False],
                     [True, True, True, True, False, False, False]])

x = x.to(device)
mask = mask.to(device)

print(x[mask])

x is tensor([-100, 0, 2, 0, 1, 0, 2, 0], device='mps:0') , what is not correct.

But when using CPU or CUDA as a device, the result is tensor([1, 0, 2, 0, 1, 0, 2, 0]), which is correct. The result is also correct when x is 1D or when it's 2D but the mask is in the CPU. The error also happens when using torch.masked_select.

Versions

PyTorch version: 2.2.2
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.4.1 (x86_64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.14 (main, Mar 21 2024, 11:24:58) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-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 Max

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.2.2
[pip3] torchaudio==2.2.2
[pip3] torchmetrics==1.3.2
[pip3] torchvision==0.17.2
[conda] numpy 1.26.4 pypi_0 pypi
[conda] torch 2.2.2 pypi_0 pypi
[conda] torchaudio 2.2.2 pypi_0 pypi
[conda] torchmetrics 1.3.2 pypi_0 pypi
[conda] torchvision 0.17.2 pypi_0 pypi

cc @frank-wei @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen

[ayush0x00]

Hey @joaquimgomez I think the issue is fixed on torch==2.3.0.dev20240212 It's working fine on my machine.

[malfet]

@ayush0x00 are you on Intel or Apple Silicon Mac? It works my M2, but it's very likely broken on Intel machines

[ayush0x00]

@ayush0x00 are you on Intel or Apple Silicon Mac? It works my M2, but it's very likely broken on Intel machines

I am using it on Apple silicon M1. Not sure about intel machines then. Any way of verifying it on intel?

[ayush0x00]

@joaquimgomez could you verify if it works when using device=torch.device("cpu")?

[joaquimgomez]

I tested it on a Linux machine with an Intel Xeon E5645 and it works.

In my M2 Max, if device=torch.device("cpu") it works @ayush0x00 . It also works if x is in MPS but mask is in the CPU. I did the tests with 2.3.0 and 2.4.0a0+git94b328e and it is still not working well when both are in MPS.

[ayush0x00]

@joaquimgomez I guess Apple has discontinued its support for Intel GPUs. Have a look at this thread. So unless you are working on Apple silicon, I think device(mps) won't work.

[joaquimgomez]

No, sorry for the misunderstanding. I'm working on an Apple Silicon, a MacBook Pro 14'' 2023 with an M2 Max.

[ayush0x00]

@joaquimgomez maybe try updating torch. It works fine on my mac, which runs M1 silicon.

[joaquimgomez]

Tried to update torch without compiling it by myself and it works now. Thanks.