import tensorflow as tf

data = tf.constant([float('nan')],dtype=tf.float32)
indices = tf.constant([1],dtype=tf.int32)
segment_ids = tf.constant([0],dtype=tf.int32)
with tf.device('/cpu:0'):
    result1 = tf.sparse.segment_mean(data,indices,segment_ids)
    print(result1)

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
Cell In[230], line 7
      5 segment_ids = tf.constant([0],dtype=tf.int32)
      6 with tf.device('/cpu:0'):
----> 7     result1 = tf.sparse.segment_mean(data,indices,segment_ids)
      8     print(result1)
      9 # with tf.device('/gpu:0'):
     10 #     result2 = tf.sparse.segment_mean(data,indices,segment_ids)
     11 #     print(result2)

File ~/anaconda3/envs/cotfuzz/lib/python3.11/site-packages/tensorflow/python/ops/math_ops.py:5106, in sparse_segment_mean_v2(data, indices, segment_ids, num_segments, name, sparse_gradient)
   5065 @tf_export("sparse.segment_mean", v1=[])
   5066 def sparse_segment_mean_v2(
   5067     data,
   (...)   5072     sparse_gradient=False,
   5073 ):
   5074   r"""Computes the mean along sparse segments of a tensor.
   5075 
   5076   Read [the section on
   (...)   5104     inferred for the last element in `segments_ids`.
   5105   """
-> 5106   return sparse_segment_mean(
   5107       data,
   5108       indices,
...
   6004 def raise_from_not_ok_status(e, name) -> NoReturn:
   6005   e.message += (" name: " + str(name if name is not None else ""))
-> 6006   raise core._status_to_exception(e) from None

InvalidArgumentError: {{function_node __wrapped__SparseSegmentMean_device_/job:localhost/replica:0/task:0/device:CPU:0}} Bad: indices[0] == 1 out of range [0, 1) [Op:SparseSegmentMean] name:

import tensorflow as tf

data = tf.constant([float('nan')],dtype=tf.float32)
indices = tf.constant([1],dtype=tf.int32)
segment_ids = tf.constant([0],dtype=tf.int32)
with tf.device('/gpu:0'):
    result2 = tf.sparse.segment_mean(data,indices,segment_ids)
    print(result2)