pytorch / pytorch/pytorch

`torch.nn.functional.PoissonNLLLoss` returns `NaN` imaginary part on CPU (vs 0 on GPU) for overflowing complex inputs

Open
#173,796 0 comments 0 reactions 0 assignees View on GitHub
module: NaNs and Infs module: nn triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

When using `torch.nn.functional.poisson_nll_loss` with `log_input=True` and `complex64` inputs that cause an overflow in the real component (e.g., `100+0j`), the CPU implementation produces a `NaN` imaginary component (`inf + nanj`), whereas the CUDA implementation correctly preserves the zero imaginary component (`inf + 0j`).

While the infinite real part is expected due to `float32` overflow, the divergence in imaginary parts suggests an inconsistency in how the two backends handle the arithmetic of `inf * 0` (likely during the underlying exponentiation).

Here is a [gist](https://colab.research.google.com/gist/jiren-the-gray/c74ac4c7b657caeecf272e914c3d888e/poisson_nll_loss.ipynb)

**Minimal reproduction**
```python
import torch
import torch.nn.functional as F

# Create a single complex element that will cause float32 overflow (e^100)
x = torch.tensor([100+0j], dtype=torch.complex64)
target = torch.tensor([0], dtype=torch.float32)

# CPU Execution
y_cpu = F.poisson_nll_loss(x, target, log_input=True)

# GPU Execution
if torch.cuda.is_available():
y_gpu = F.poisson_nll_loss(x.cuda(), target.cuda(), log_input=True)

print(f"CPU: {y_cpu} (Imag is NaN: {y_cpu.imag.isnan().item()})")
print(f"GPU: {y_gpu.cpu()} (Imag is NaN: {y_gpu.imag.isnan().item()})")
else:
print("CUDA not available for comparison.")
```

**Output**
```
CPU: (inf+nanj) (Imag is NaN: True)
GPU: (inf+0j) (Imag is NaN: False)
```

**Expected behavior**
CPU and GPU should produce consistent results. Ideally, the CPU implementation should handle the `inf * sin(0)` case gracefully to return `0j` rather than `nanj`, matching the CUDA behavior.

### Versions

Collecting environment information...
PyTorch version: 2.10.0+cu128
Is debug build: False
CUDA used to build PyTorch: 12.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04.2) 11.4.0
Clang version: Could not collect
CMake version: version 3.31.10
Libc version: glibc-2.35

Python version: 3.12.12 (main, Oct 10 2025, 08:52:57) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.6.105+-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.5.82
CUDA_MODULE_LOADING set to:
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 550.54.15
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.2.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.2.1
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 2
On-line CPU(s) list: 0,1
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU @ 2.00GHz
CPU family: 6
Model: 85
Thread(s) per core: 2
Core(s) per socket: 1
Socket(s): 1
Stepping: 3
BogoMIPS: 4000.28
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat md_clear arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 32 KiB (1 instance)
L1i cache: 32 KiB (1 instance)
L2 cache: 1 MiB (1 instance)
L3 cache: 38.5 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0,1
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Vulnerable
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Mitigation; PTE Inversion
Vulnerability Mds: Vulnerable; SMT Host state unknown
Vulnerability Meltdown: Vulnerable
Vulnerability Mmio stale data: Vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Vulnerable
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2: Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Vulnerable
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Vulnerable

Versions of relevant libraries:
[pip3] intel-cmplr-lib-ur==2025.3.1
[pip3] intel-openmp==2025.3.1
[pip3] mkl==2025.3.0
[pip3] numpy==2.0.2
[pip3] nvidia-cublas-cu12==12.8.4.1
[pip3] nvidia-cuda-cupti-cu12==12.8.90
[pip3] nvidia-cuda-nvrtc-cu12==12.8.93
[pip3] nvidia-cuda-runtime-cu12==12.8.90
[pip3] nvidia-cudnn-cu12==9.10.2.21
[pip3] nvidia-cufft-cu12==11.3.3.83
[pip3] nvidia-curand-cu12==10.3.9.90
[pip3] nvidia-cusolver-cu12==11.7.3.90
[pip3] nvidia-cusparse-cu12==12.5.8.93
[pip3] nvidia-cusparselt-cu12==0.7.1
[pip3] nvidia-nccl-cu12==2.27.5
[pip3] nvidia-nvjitlink-cu12==12.8.93
[pip3] nvidia-nvtx-cu12==12.8.90
[pip3] nvtx==0.2.14
[pip3] onemkl-license==2025.3.0
[pip3] optree==0.18.0
[pip3] tbb==2022.3.0
[pip3] tcmlib==1.4.1
[pip3] torch==2.10.0
[pip3] torchao==0.10.0
[pip3] torchaudio==2.9.0+cu126
[pip3] torchdata==0.11.0
[pip3] torchsummary==1.5.1
[pip3] torchtune==0.6.1
[pip3] torchvision==0.24.0+cu126
[pip3] triton==3.6.0
[pip3] umf==1.0.2
[conda] Could not collect

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.