GPU tf.math.zeta returns NaN for finite s and infinite q #121501

#!/usr/bin/env python3
import os
import sys

os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0")

import numpy as np
import tensorflow as tf

s = np.array([2.0], dtype=np.float32)
q = np.array([np.inf], dtype=np.float32)
correct = np.array([0.0], dtype=np.float32)

gpus = tf.config.list_logical_devices("GPU")
if not gpus:
    print("missing target: TensorFlow GPU is not available")
    sys.exit(2)

with tf.device("/GPU:0"):
    wrong = tf.math.zeta(tf.constant(s), tf.constant(q)).numpy()

print("wrong:", wrong.tolist())
print("correct:", correct.tolist())

if not np.array_equal(wrong, correct, equal_nan=True):
    sys.exit(0)

sys.exit(1)

wrong: [nan]
correct: [0.0]