tf.math.angle(nan) returns inconsistent results on dtype float64 / complex128.

[QuantumCoder4]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

binary

TensorFlow version

tf 2.15.0

Custom code

Yes

OS platform and distribution

Linux-5.14.0-362.18.1.el9_3.x86_64-x86_64-with-glibc2.34

Mobile device

AlmaLinux 9

Python version

3.9

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

tf.math.angle gives inconsistent outputs on NaN value: if tensor dtype=complex128 then NaN, while for dtype=float64 it is 0.
Although doc states that for any real input the output will be zero, I would like to list some points to argue.

1. It does not make sense for ending up with a regular value for calculating NaN.
2. The nature of NaN should not change, regardless of the dtype of the tensor it's staying in.
3. Such behavior may let NaN error to escape, causing trouble in debugging.

Expected behavior: the result is NaN when input is NaN.

Standalone code to reproduce the issue

input = tf.constant([np.nan], dtype=tf.float64)
out = tf.math.angle(input)  # tf.Tensor([0.], shape=(1,), dtype=float64)

input = tf.constant([np.nan], dtype=tf.complex128)
out = tf.math.angle(input)  # tf.Tensor([nan], shape=(1,), dtype=float64)

Relevant log output

No response

[NBCBM]

import tensorflow as tf
import numpy as np

def angle_with_nan_handling(input_tensor):
# Check if input_tensor contains NaN values
contains_nan = tf.reduce_any(tf.math.is_nan(input_tensor))

# If input_tensor contains NaN values, return NaN
if contains_nan:
    return tf.fill(tf.shape(input_tensor), tf.constant(np.nan, dtype=tf.float64))
else:
    # Calculate angle for non-NaN values
    return tf.math.angle(input_tensor)

Test with float64 input

input_float64 = tf.constant([np.nan], dtype=tf.float64)
out_float64 = angle_with_nan_handling(input_float64)
print("Output for float64 input:", out_float64.numpy()) # Output: [nan]

Test with complex128 input

input_complex128 = tf.constant([np.nan], dtype=tf.complex128)
out_complex128 = angle_with_nan_handling(input_complex128)
print("Output for complex128 input:", out_complex128.numpy()) # Output: [nan]

[sushreebarsa]

@SuryanarayanaY I was able to replicate the issue here.
Thank you!