Fix keras.ops.quantile implementation for floating point inputs that …

…are not tf.float32. (keras-team#20438)
wang-xianghao · Nov 20, 2024 · 233574f · 233574f
1 parent 983b377
commit 233574f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
@@ -1799,7 +1799,8 @@ def _get_indices(method):
         nan_batch_members = tf.reshape(
             nan_batch_members, shape=right_rank_matched_shape
         )
-        gathered_y = tf.where(nan_batch_members, float("NaN"), gathered_y)
+        nan_value = tf.constant(float("NaN"), dtype=x.dtype)
+        gathered_y = tf.where(nan_batch_members, nan_value, gathered_y)
 
     # Expand dimensions if requested
     if keepdims: