[MPS] torch.fft.hfft aborts on 5D complex64 input with dim=0
- 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
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