How can I validate the segment_ids argument of tf.unsorted_segments_max?

Question

In TensorFlow 1.12 you can use tf.math.unsorted_segment_sum to compute the maximum along segments of a tensor.

One of the arguments:

segment_ids: A Tensor. Must be one of the following types: int32, int64. A tensor whose shape is a prefix of data.shape.END } out_arg { name: "output" description: << END Has same shape as data, except for the first segment_ids.rank dimensions, which are replaced with a single dimension which has size num_segments.

Q1: I don't understand what this means. I suppose segment_ids should be [0,1,2,3,...], and serial number starts with 0.

Then I tried to verify my assumption by testing with different values for segment_ids:

print(sess.run(tf.unsorted_segment_max(tf.constant([0.1, 0.2, 0.3, 0.3, 0.4]),
                                       tf.constant([2, 0, 1, 1, 0]), 3)))
#[0.4 0.3 0.1], correct
print(sess.run(tf.unsorted_segment_max(tf.constant([0.1, 0.2, 0.3, 0.3, 0.4]),
                                       tf.constant([3, 0, 1, 1, 0]), 3)))
#[ 4.0000001e-01  3.0000001e-01 -3.4028235e+38], number 2 worked, but 3 didnt
print(sess.run(tf.unsorted_segment_max(tf.constant([0.1, 0.2, 0.3, 0.3, 0.4]),
                                       tf.constant([3, 0, 1, 1, 0]), 4)))
#[ 4.0000001e-01  3.0000001e-01 -3.4028235e+38  1.0000000e-01], num_segments didnt help
print(sess.run(tf.unsorted_segment_sum(tf.constant([0.1, 0.2, 0.3, 0.3, 0.4]),
                                       tf.constant([3, 0, 1, 1, 0]), 4)))
#[0.6 0.6 0.  0.1], while num_segments worked on sum.

From Ex3 and Ex4, I notice that segment_ids in tf.unsorted_segment_max() got extra limitation other than tf.unsorted_segment_sum(), and I suppose it is consecutive from 0. Ex3 shall have [0.4 0.3 0. 0.1]

From Ex1 and Ex2, I notice that Ex1 result is right, while Ex2 shall have [0.4 0.3 0.] and there is no warning or error ocured if input segment_ids is invalid.

Q2: How can I check my segment_ids to see if it fits the requirements?

Q3: If my segment_ids doesn't fit the requirement, like [3, 0, 1, 1, 0], what can I do to make unsorted_segments_max work?

Answer 1

A1: For the one-dimensional case the output tensor will have as many elements as specified by num_segments. Each element i will be the result of the operation (e.g. max, sum, ...) being applied to all elements in data, denoted by position j, for which segment_ids[j] == i. The segment_ids can be any numbers and don't necessarily need to start at 0. Nevertheless the output will contain num_segments elements and if for any output element i no segment_ids[j] == i is found then a specific default value is used (this value is different for each of the operations). For higher dimensional tensors the first x dimensions of data, as specified by the rank of segment_ids, will be replaced by such a single dimension containing the various segments.

Your examples

You didn't seem to run your examples on tensorflow >= 1.12.0 because the second example should have raised an exception: InvalidArgumentError (see above for traceback): segment_ids[0] = 3 is out of range [0, 3). The fact that it is ignored is probably the behavior in older versions (please check).

From the documentation of tf.math.unsorted_segment_max:

If the maximum is empty for a given segment ID i, it outputs the smallest possible value for the specific numeric type, output[i] = numeric_limits<T>::lowest().

1. You specify 3 segments and so the output has 3 elements. The first element, i == 0, corresponds to locations j == [1, 4] (as specified in segment_ids) and hence data[j] == [0.2, 0.4]. Taking the max(data[j]) yields 0.4. Similarly for the next segment, i == 1, we have j == [2, 3] and data[j] = [0.3, 0.3] => max(data[j]) = 0.3. For the last segment, i == 2, there is only a single element at position j == 0 and hence the result is 0.1.
2. Now you still specify 3 clusters and so the output will have 3 elements. The segment_id 3 seems to be ignored because it doesn't fit the number of segments (this should raise an error on the latest version of tensorflow though). Regarding the last element in the output, which corresponds to the segment i == 2, no corresponding segment_id is found and hence the result is the smallest possible value for the specific numeric type; this is what you observe: -3.4028235e+38.
3. Similar to 2) with the difference that now the segment_id 3 becomes effective; this corresponds to the last element in the output which is thus 0.1. The second to last element, corresponding to segment_id 2 has still no matching ids, and is hence filled by the default value.
4. For tf.math.unsorted_segment_sum the default value for empty segments is different: If the sum is empty for a given segment ID i, output[i] = 0. Hence segment 0 is the sum of items 0.2 + 0.4 == 0.6, for segment 1 it is 0.3 + 0.3 == 0.6, for segment 2 no corresponding id can be found and thus it is 0 by default and for segment 3 there is only a single element 0.1.

A2: The criteria for segment_ids are that it "prefixes" the shape of data, i.e. segment_ids.shape == data.shape[:len(segment_ids.shape)] and that max(segment_ids) < num_segments (on the latest version of tensorflow). The corresponding dimensions will then be reduced to a single one, containing as many elements as specified by num_segments while the values in segments_ids specify the corresponding segments in the output.

A3: In case the segment_ids argument doesn't meet the requirements an InvalidArgumentError will be raised. You'll have to fix the corresponding tensor in order to make the operation work.