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]

from tensorflow.

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

from tensorflow.