[Bug]: Softmax produces extra NaN with mixed inf input on GPU devices

[zhihaoxu1325]

OpenVINO Version

2025.4.0

Operating System

Ubuntu 20.04 (LTS)

Device used for inference

GPU

Framework

None

Model used

No response

Issue description

The GPU implementation of Softmax operation produces incorrect results when the input contains a mix of inf and finite values. Specifically, positions that should output 0.0 incorrectly output NaN.

Expected Behavior (IEEE 754)
For input [inf, 1.0, 2.0]:

softmax([inf, 1.0, 2.0]) = [exp(inf), exp(1), exp(2)] / sum
                        = [inf, e, e²] / inf
                        = [inf/inf, e/inf, e²/inf]
                        = [nan, 0.0, 0.0]

The first element becomes NaN (inf/inf is undefined), but other elements should be 0.0 (finite/inf = 0).

Actual Behavior

CPU: [[nan, 0.0, 0.0]]  ← Correct
GPU: [[nan, nan, nan]]  ← Wrong: all elements become NaN

Step-by-step reproduction

Please run the script to reproduce the issues:

#!/usr/bin/env python3
"""
Bug 2: Softmax produces extra NaN with mixed inf input on GPU
"""

import numpy as np
import openvino as ov
import openvino.runtime.opset13 as ops
from openvino.runtime import Model

def test_softmax_inf_bug():
    print("=" * 70)
    print("Bug: [GPU] Softmax produces extra NaN with mixed inf")
    print("=" * 70)
    
    core = ov.Core()
    print(f"OpenVINO version: {ov.__version__}")
    
   
    x = np.array([[np.inf, 1.0, 2.0]], dtype=np.float32)
    
    print(f"\nInput: {x}")
    print(f"\nExpected behavior:")
    print(f"  softmax([inf, 1, 2]) = [exp(inf), exp(1), exp(2)] / sum")
    print(f"  = [inf, e, e^2] / inf = [inf/inf, e/inf, e^2/inf] = [nan, 0, 0]")
    
    param = ops.parameter(x.shape, np.float32)
    softmax = ops.softmax(param, axis=-1)
    model = Model([softmax], [param])
    
    cpu_result = np.array(core.compile_model(model, "CPU")([x])[0])
    gpu_result = np.array(core.compile_model(model, "GPU")([x])[0])
    
    print(f"\nResults:")
    print(f"  CPU: {cpu_result}")
    print(f"  GPU: {gpu_result}")
    
    cpu_nan_count = np.sum(np.isnan(cpu_result))
    gpu_nan_count = np.sum(np.isnan(gpu_result))
    
    if cpu_nan_count != gpu_nan_count:
        print(f"\n*** BUG CONFIRMED ***")
        print(f"  CPU NaN count: {cpu_nan_count}")
        print(f"  GPU NaN count: {gpu_nan_count}")
        print(f"  GPU Softmax produces {gpu_nan_count - cpu_nan_count} extra NaN values")
    
    
    print(f"\n" + "-" * 70)
    print("Additional test cases:")
    
    test_cases = [
        np.array([[np.inf, -np.inf, 1.0]], dtype=np.float32),
        np.array([[np.inf, np.inf, 1.0]], dtype=np.float32),
        np.array([[-np.inf, 1.0, 2.0]], dtype=np.float32), 
    ]
    
    for tc in test_cases:
        cpu_r = np.array(core.compile_model(model, "CPU")([tc])[0])
        gpu_r = np.array(core.compile_model(model, "GPU")([tc])[0])
        cpu_nan = np.sum(np.isnan(cpu_r))
        gpu_nan = np.sum(np.isnan(gpu_r))
        status = "OK" if cpu_nan == gpu_nan else "MISMATCH"
        print(f"  {tc.flatten()} → CPU:{cpu_r.flatten()}, GPU:{gpu_r.flatten()} [{status}]")

if __name__ == "__main__":
    test_softmax_inf_bug()

Relevant log output

OpenVINO version: 2025.4.0-20297-5f5229a2009

Input: [[inf  1.  2.]]

Expected behavior:
  softmax([inf, 1, 2]) = [exp(inf), exp(1), exp(2)] / sum
  = [inf, e, e^2] / inf = [inf/inf, e/inf, e^2/inf] = [nan, 0, 0]

Results:
  CPU: [[nan  0.  0.]]
  GPU: [[nan nan nan]]

*** BUG CONFIRMED ***
  CPU NaN count: 1
  GPU NaN count: 3
  GPU Softmax produces 2 extra NaN values

----------------------------------------------------------------------
Additional test cases:
  [ inf -inf   1.] → CPU:[nan  0.  0.], GPU:[nan nan nan] [MISMATCH]
  [inf inf  1.] → CPU:[nan nan  0.], GPU:[nan nan nan] [MISMATCH]
  [-inf   1.   2.] → CPU:[0.         0.26894143 0.73105854], GPU:[0.         0.26894143 0.7310586 ] [OK]

Issue submission checklist

• I'm reporting an issue. It's not a question.
• I checked the problem with the documentation, FAQ, open issues, Stack Overflow, etc., and have not found a solution.
• There is reproducer code and related data files such as images, videos, models, etc.

[AyusKumarPathak]

Hi maintainers 👋

I’d like to work on this Softmax NaN discrepancy on GPU.

Plan of action:

1. Reproduce the issue locally on GPU using the repro script in the description.
2. Investigate the GPU kernel implementation of Softmax and trace how inf values are handled.
3. Update the logic to correctly implement IEEE 754 behavior such that only the mathematically correct NaN remains and other values evaluate to 0.0.
4. Add unit tests verifying correct behavior with mixed inf/finite inputs.
5. Ensure tests pass on both CPU and GPU.

I expect to open a draft PR with tests and diagnostics in the next 48–72 hours.

Please assign this issue to me if this approach is acceptable.

— AyusKumarPathak

[zhihaoxu1325]

Hi maintainers 👋

I’d like to work on this Softmax NaN discrepancy on GPU.

Plan of action:

1. Reproduce the issue locally on GPU using the repro script in the description.
2. Investigate the GPU kernel implementation of Softmax and trace how inf values are handled.
3. Update the logic to correctly implement IEEE 754 behavior such that only the mathematically correct NaN remains and other values evaluate to 0.0.
4. Add unit tests verifying correct behavior with mixed inf/finite inputs.
5. Ensure tests pass on both CPU and GPU.

I expect to open a draft PR with tests and diagnostics in the next 48–72 hours.

Please assign this issue to me if this approach is acceptable.

— AyusKumarPathak

hi @AyusKumarPathak you can comment .take in this issues so that it will be assigned to you.

[AyusKumarPathak]

.take