Tensor.abs() gives incorrect results on Complex64 when using MPS

[stevenryoung]

🐛 Describe the bug

The abs() function gives incorrect results when using Complex64 tensors on an MPS device. It appears to be some kind of indexing issue per the example below.

import torch

with torch.device("cpu"):
    print(torch.ones((2,), dtype=torch.complex64).abs())
    

with torch.device("mps"):
    print(torch.ones((2,), dtype=torch.complex64).abs())
    print(torch.tensor([1.0 + 0.0j, 0.0 + 10.0j, 100.0 + 0.0j, 1000.0 + 0.0j]).abs())

Output:

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

Versions

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

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

Python version: 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:34:54) [Clang 16.0.6 ] (64-bit runtime)
Python platform: macOS-14.4.1-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.dev20240428
[pip3] torchaudio==2.2.0.dev20240428
[pip3] torchvision==0.19.0.dev20240428
[conda] numpy 1.26.4 py311h7125741_0 conda-forge
[conda] pytorch 2.4.0.dev20240428 py3.11_0 pytorch-nightly
[conda] torchaudio 2.2.0.dev20240428 py311_cpu pytorch-nightly
[conda] torchvision 0.19.0.dev20240428 py311_cpu pytorch-nightly

cc @ezyang @anjali411 @dylanbespalko @mruberry @lezcano @nikitaved @amjames @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen

[stevenryoung]

For any user needing a quick (probably unsafe) fix, the following worked for my use case:

def patch_abs(t):
    if (t.device == torch.device("mps:0")) and (t.dtype == torch.complex64):
        return torch.sqrt(torch.pow(torch.real(t), 2) + torch.pow(torch.imag(t), 2) + 1e-12)
    return torch.abs(t)

torch.Tensor.abs = patch_abs

with torch.device("cpu"):
    print(torch.ones((2,), dtype=torch.complex64).abs())
    
with torch.device("mps"):
    print(torch.ones((2,), dtype=torch.complex64).abs())
    print(torch.tensor([1.0 + 0.0j, 0.0 + 10.0j, 100.0 + 0.0j, 1000.0 + 0.0j]).abs())

Output:

tensor([1., 1.])
tensor([1., 1.], device='mps:0')
tensor([   1.,   10.,  100., 1000.], device='mps:0')

[gambiTarun]

Hi @bdhirsh !
First time contributor here. Could I work on this bug?

[malfet]

Grabbing for myself, wish there was a proper doc for https://developer.apple.com/documentation/metalperformanceshadersgraph/mpsgraph/3564540-absolutewithtensor?language=objc

[malfet]

@gambiTarun sorry, missed your comment. How familiar are you with ObjC/MPS framework?

[gambiTarun]

Hi @malfet, no worries! I’m not familiar with ObjC/MPS, but I’m willing to learn. If you could point me in the right direction for resolving this bug, or suggest some resources, I'd appreciate it!

[jhavukainen]

Hi @malfet! Wanted to add some context as I did have a brief discussion with some of the MPSGraph people regarding this issue and that particular op. My current hypothesis on what happens:

• If MPSGraph abs op gets passed an op that is data type Complex64, it will return a tensor of the same type. However if I print the type of the output from the CPU abs() (and the MPS abs for that matter) using the above snippet provided in the error it seems to be in float32. So probably the complex parts get dropped due to this conversion not being handled explicitly/correctly in the MPS abs implementation and we end up just interpreting the complex tensor as a float32 tensor. This would make complex 1+0i look like a float2 [1,0] tensor.
• If that's the case, we should check for complex data type in the abs op and explicitly take the real part of the complex output of the graph op absoluteWithTensor. There should be an API

-(MPSGraphTensor *) realPartOfTensor:(MPSGraphTensor *) tensor name:(NSString * _Nullable) name
that should do the trick.

Unfortunately I'm drowning in other work this week so I'm more than happy to let anyone take a shot at this if they'd like.