pytorch / pytorch/pytorch

[MPS] torch.fft.hfft aborts on 5D complex64 input with dim=0

Open
#188,864 1 comment 0 reactions 0 assignees View on GitHub
bot-triaged module: complex module: crash module: error checking module: fft module: mps topic: fuzzer triaged
Dominant language
Python
Stars
103k
Forks
29.6k
PR merge metrics
PR metrics pending

Description

### πŸ› Describe the bug

## Bug: torch.fft.hfft aborts Python process on MPS for 5D complex64 input with dim=0

### Description

`torch.fft.hfft` with a 5D `complex64` input tensor on MPS and `dim=0` aborts the
Python process (SIGABRT, exit code 134) instead of returning a result or raising a
catchable Python exception. The same call succeeds on CPU.

The failure originates in MPSGraph's `mps.hermitean_to_real_fft` op, which rejects
`axis=0` for rank-5 tensors: the Metal FFT transform is only supported on the last
four dimensions (`axis` must be > 0). PyTorch's MPS backend does not validate this
constraint before graph compilation, so the invalid graph is built and MPSGraph
asserts during module verification, killing the process.

The fix should either (a) implement rank-5 / `dim=0` support in the MPS FFT path to
match CPU, or (b) reject the call at the PyTorch dispatch boundary with a clear
`RuntimeError` before MPSGraph compilation.

### To Reproduce

```python
import torch

assert torch.backends.mps.is_available(), "MPS not available β€” run on Apple Silicon"
device = torch.device("mps")

x = torch.zeros(2, 2, 2, 3, 5, dtype=torch.complex64, device=device)
torch.fft.hfft(x, dim=0)
```

### Expected behavior

Return the correct `float32` result on MPS, matching CPU:

```python
import torch

x = torch.zeros(2, 2, 2, 3, 5, dtype=torch.complex64)
y = torch.fft.hfft(x, dim=0)
# y.shape == (9, 2, 2, 3, 5) dtype=float32
```

Or, if the MPS FFT kernel cannot support this axis/rank combination, raise a clear
`RuntimeError` instead of aborting the process.

### Actual behavior

Process abort (SIGABRT). The process dies with messages like:

```
'mps.hermitean_to_real_fft' op invalid axis: 0 for rank: 5. Transform supported
only on the last four dimensions, ie. axis must be larger than 0.

MPSGraphExecutable.mm:1484: failed assertion `original module failed verification'
```

### Additional context

CPU correctly handles the same call:

```python
import torch

x = torch.zeros(2, 2, 2, 3, 5, dtype=torch.complex64)
y = torch.fft.hfft(x, dim=0)
# succeeds: y.shape == (9, 2, 2, 3, 5)
```

MPS succeeds when transforming along a supported axis (e.g. `dim=-1`):

```python
import torch

x = torch.zeros(2, 2, 2, 3, 5, dtype=torch.complex64, device="mps")
y = torch.fft.hfft(x, dim=-1) # OK
```

**Related issues:** This is the same class of MPSGraph hard-assert failure seen when
invalid FFT graphs are compiled instead of being rejected at the PyTorch boundary.
The underlying constraint is that `mps.hermitean_to_real_fft` only supports transforms
on the last four tensor dimensions; any 5D input with `dim=0` (or `dim` mapping to
axis 0) triggers the abort.

### Environment

- PyTorch version: 2.12.0
- OS: macOS arm64 (Apple Silicon)
- Backend: MPS

### Affected scope

| Field | Value |
|-----------------|--------------------------------------------------------------|
| Op | `torch.fft.hfft` |
| Device | `mps` (Apple Silicon) |
| Dtype | `torch.complex64` |
| Trigger | `ndim == 5` and `dim=0` (axis 0 unsupported by MPSGraph FFT) |
| Not affected | same op on CPU; MPS with `dim` on last four axes (e.g. `-1`) |
| Likely affected | `torch.fft.hfft2`, `torch.fft.hfftn` on MPS for high-rank inputs with early-axis `dim` |
"""
```
import torch

SHAPE = (2, 2, 2, 3, 5)
DIM_OK = -1
DIM_BUG = 0

print(f"torch version: {torch.__version__}")
assert torch.backends.mps.is_available(), "MPS not available β€” run on Apple Silicon"

device = torch.device("mps")
```

### Baseline: hfft on last axis should succeed on MPS ───────────────────────
```
print("[baseline] complex64 hfft on MPS (dim=-1):")
x_ok = torch.zeros(SHAPE, dtype=torch.complex64, device=device)
y_ok = torch.fft.hfft(x_ok, dim=DIM_OK)
print(f" shape={SHAPE} dim={DIM_OK} -> output shape={tuple(y_ok.shape)} dtype={y_ok.dtype}")
```
### CPU reference: dim=0 should succeed ─────────────────────────────────────
```print("\n[ref] complex64 hfft on CPU (dim=0):")
x_cpu = torch.zeros(SHAPE, dtype=torch.complex64)
y_cpu = torch.fft.hfft(x_cpu, dim=DIM_BUG)
print(f" shape={SHAPE} dim={DIM_BUG} -> output shape={tuple(y_cpu.shape)} dtype={y_cpu.dtype}")
```
### Failing case: dim=0 on MPS aborts the process (fuzz example) ────────────
```
print("\n[bug] complex64 hfft on MPS (dim=0) β€” expect process abort:")
x_bug = torch.zeros(SHAPE, dtype=torch.complex64, device=device)
print(f" shape={SHAPE} dim={DIM_BUG} dtype={x_bug.dtype}")
torch.fft.hfft(x_bug, dim=DIM_BUG)
print(" UNEXPECTED ok")
```

### Versions

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

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

Python version: 3.14.5 (main, May 10 2026, 19:20:57) [Clang 22.1.3 ] (64-bit runtime)
Python platform: macOS-26.5-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 M5

Versions of relevant libraries:
[pip3] numpy==2.4.6
[pip3] torch==2.12.0
[conda] Could not collect
```

cc @malfet @ezyang @anjali411 @dylanbespalko @mruberry @nikitaved @amjames @kulinseth @DenisVieriu97 @jhavukainen @aditvenk @Isalia20

Contributor guide

Open the contributing guide

Research direction

Start by running the supplied MPS reproduction and compare it with the CPU call and the working dim=-1 case. Trace the MPS FFT path for torch.fft.hfft and its dispatch boundary; done means rank-5 dim=0 no longer aborts, either by matching CPU or by raising a clear RuntimeError before MPSGraph compilation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.