tf.math.multiply for Complex inf/nan Gives Inconsistent Results on CPU for Tensors with >= 2 Elements #98408
Description
Issue type
Bug
Have you reproduced the bug with TensorFlow Nightly?
Yes
Source
source
TensorFlow version
2.17
Custom code
Yes
OS platform and distribution
No response
Mobile device
No response
Python version
No response
Bazel version
No response
GCC/compiler version
No response
CUDA/cuDNN version
No response
GPU model and memory
No response
Current behavior?
tf.math.multiply exhibits a size-dependent inconsistency when multiplying complex numbers involving inf and nan. The GPU's behavior diverges from the CPU's behavior only when the input tensor contains two or more elements.
Standalone code to reproduce the issue
For a single-element tensor:
import tensorflow as tf
import numpy as np
import os
input_a = tf.constant([
complex(np.inf, np.inf)
], dtype=tf.complex128)
input_b = tf.constant([
complex(1, np.nan)
], dtype=tf.complex128)
with tf.device('/CPU:0'):
cpu_out = tf.math.multiply(input_a, input_b)
print(f"CPU output:\n{cpu_out.numpy()}")
with tf.device('/GPU:0'):
gpu_out = tf.math.multiply(input_a, input_b)
print(f"\nGPU output:\n{gpu_out.numpy()}")output:
CPU output:
[inf+infj]
GPU output:
[inf+infj]
But for a tensor with >= 2 elements:
import tensorflow as tf
import numpy as np
import os
input_a = tf.constant([
complex(np.inf, np.inf),complex(np.inf, np.inf)
], dtype=tf.complex128)
input_b = tf.constant([
complex(1, np.nan),complex(1, np.nan)
], dtype=tf.complex128)
with tf.device('/CPU:0'):
cpu_out = tf.math.multiply(input_a, input_b)
print(f"CPU output:\n{cpu_out.numpy()}")
with tf.device('/GPU:0'):
gpu_out = tf.math.multiply(input_a, input_b)
print(f"\nGPU output:\n{gpu_out.numpy()}")output:
CPU output:
[nan+nanj nan+nanj]
GPU output:
[inf+infj inf+infj]