def max_reducer(neighbor_msg, node_index, num_nodes=None):
if num_nodes is None:
num_nodes = tf.reduce_max(node_index) + 1
# return tf.math.unsorted_segment_max(neighbor_msg, node_index, num_segments=num_nodes)
# max_x = tf.math.unsorted_segment_max(x, node_graph_index, num_segments=num_graphs)
max_neighbor_msg = segment_op_with_pad(tf.math.segment_max, neighbor_msg, node_index, num_segments=num_nodes)
return max_neighbor_msg
def max_reducer(neighbor_msg, node_index, num_nodes=None):
if num_nodes is None:
num_nodes = tf.reduce_max(node_index) + 1
max_neighbor_msg = segment_op_with_pad(tf.math.segment_max,
neighbor_msg,
node_index,
num_segments=num_nodes)
return max_neighbor_msg
def min_pool(x, node_graph_index, num_graphs=None):
if num_graphs is None:
num_graphs = tf.reduce_max(node_graph_index) + 1
# min_x = tf.math.unsorted_segment_min(x, node_graph_index, num_segments=num_graphs)
min_x = segment_op_with_pad(tf.math.segment_min,
x,
node_graph_index,
num_segments=num_graphs)
return min_x