pytorch - 💡(How to fix) Fix [MPS] Boolean indexing returns inconsistent tensor sizes on repeated calls [1 participants]

🐛 Describe the bug

On the MPS backend, boolean indexing with the same 1D tensor and the same 1D boolean mask can return tensors with inconsistent sizes across repeated calls.

x[mask].numel() is expected to always equal mask.sum(), but on MPS it sometimes returns a different size. The same code works correctly on CPU.

Minimal repro:

import torch
import platform

print("torch:", torch.__version__)
print("python:", platform.python_version())
print("platform:", platform.platform())
print("mps available:", torch.backends.mps.is_available())

device = "mps"
torch.manual_seed(0)

n = 11_842_339
m = 6_054_691

x = torch.randn(n, device=device)

mask = torch.zeros(n, dtype=torch.bool, device=device)
mask[:m] = True
mask = mask[torch.randperm(n, device=device)]

expected = mask.sum().item()
print("expected:", expected)

for step in range(10_000):
    a = x[mask]
    b = x[mask]

    if a.numel() != expected or b.numel() != expected or a.numel() != b.numel():
        print("BUG FOUND")
        print("step:", step)
        print("expected:", expected)
        print("a numel:", a.numel())
        print("b numel:", b.numel())
        break
else:
    print("No bug reproduced")

Observed result on MPS:

step: 228
expected: 6054691
actual: 7123614
y shape: torch.Size([7123614])

platform: macOS-26.4.1-arm64-arm-64bit-Mach-O

Versions

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

OS: macOS 26.4.1 (arm64) GCC version: Could not collect Clang version: 21.0.0 (clang-2100.0.123.102) CMake version: version 4.3.0 Libc version: N/A

Python version: 3.13.9 | packaged by Anaconda, Inc. | (main, Oct 21 2025, 19:11:29) [Clang 20.1.8 ] (64-bit runtime) Python platform: macOS-26.4.1-arm64-arm-64bit-Mach-O 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 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Apple M4

Versions of relevant libraries: [pip3] numpy==2.4.1 [pip3] torch==2.9.1 [pip3] torchaudio==2.9.1 [pip3] torchvision==0.24.1 [conda] numpy 2.3.5 pypi_0 pypi [conda] numpy-base 2.4.1 py313h23175f9_0 [conda] torch 2.9.1 pypi_0 pypi [conda] torchaudio 2.9.1 pypi_0 pypi [conda] torchvision 0.24.1 pypi_0 pypi

cc @kulinseth @malfet @DenisVieriu97 @jhavukainen @aditvenk