Inconsistent reduction code in graph convolution

I was browsing through the graph convolution code and saw the reduction code (including the not yet published max reduction option), which currently reads like this:

if reduction == "weighted":
  edge_features_weighted = edge_features * tf.expand_dims(
      adjacency.values, -1)
  features = tf.math.unsorted_segment_sum(
      data=edge_features_weighted,
      segment_ids=adjacency_ind_0,
      num_segments=tf.shape(input=x_flat)[0])
elif reduction == "max":
  features = tf.math.segment_max(data=edge_features,
                                 segment_ids=adjacency_ind_0)
else:
  raise ValueError("The reduction method must be 'weighted' or 'max'")

As can be seen, the weighted reduction method uses tf.math.unsorted_segment_sum, while the max reduction method uses tf.math.segment_max. This does not seem to make much sense, since both are using the same adjacency_ind_0 as segment_ids. If the (possibly reshaped) sparse tensor adjacency is assumed to be ordered (and, according to the documentation of tf.sparse.SparseTensor, "most ops assume correct ordering", and I'm not fully sure if convert_to_block_diag_2d maintains the order, but I think it does), then tf.math.segment_sum should be preferred for the weighted reduction. Conversely, if unordered sparse tensors are meant to be supported, then tf.math.unsorted_segment_max should be used for the max reduction.

tensorflow / graphics

Inconsistent reduction code in graph convolution #281