def call(self, inputs):
if self.data_mode == 'graph':
X = inputs[0]
I = inputs[1]
else:
X = inputs
if self.data_mode == 'graph':
return tf.segment_sum(X, I)
else:
return K.sum(X, axis=-2, keepdims=(self.data_mode == 'single'))
def top_k(scores, I, ratio, top_k_var):
"""
Returns indices to get the top K values in `scores` segment-wise, with
segments defined by I. K is not fixed, but it is defined as a ratio of the
number of elements in each segment.
:param scores: a rank 1 tensor with scores;
:param I: a rank 1 tensor with segment IDs;
:param ratio: float, ratio of elements to keep for each segment;
:param top_k_var: a tf.Variable without shape validation (e.g.,
`tf.Variable(0.0, validate_shape=False)`);
:return: a rank 1 tensor containing the indices to get the top K values of
each segment in `scores`.
"""
num_nodes = tf.segment_sum(tf.ones_like(I),
I) # Number of nodes in each graph
cumsum = tf.cumsum(num_nodes) # Cumulative number of nodes (A, A+B, A+B+C)
cumsum_start = cumsum - num_nodes # Start index of each graph
n_graphs = tf.shape(num_nodes)[0] # Number of graphs in batch
max_n_nodes = tf.reduce_max(num_nodes) # Order of biggest graph in batch
batch_n_nodes = tf.shape(I)[0] # Number of overall nodes in batch
to_keep = tf.ceil(ratio * tf.cast(num_nodes, tf.float32))
to_keep = tf.cast(to_keep, tf.int32) # Nodes to keep in each graph
index = tf.range(batch_n_nodes)
index = (index - tf.gather(cumsum_start, I)) + (I * max_n_nodes)
y_min = tf.reduce_min(scores)
dense_y = tf.ones((n_graphs * max_n_nodes, ))
dense_y = dense_y * tf.cast(
y_min - 1, tf.float32
) # subtract 1 to ensure that filler values do not get picked
dense_y = tf.assign(
top_k_var, dense_y, validate_shape=False
) # top_k_var is a variable with unknown shape defined in the elsewhere
dense_y = tf.scatter_update(dense_y, index, scores)
dense_y = tf.reshape(dense_y, (n_graphs, max_n_nodes))
perm = tf.argsort(dense_y, direction='DESCENDING')
perm = perm + cumsum_start[:, None]
perm = tf.reshape(perm, (-1, ))
to_rep = tf.tile(tf.constant([1., 0.]), (n_graphs, ))
rep_times = tf.reshape(
tf.concat((to_keep[:, None], (max_n_nodes - to_keep)[:, None]), -1),
(-1, ))
mask = tf_repeat_1d(to_rep, rep_times)
perm = tf.boolean_mask(perm, mask)
return perm
def call(self, inputs):
if self.data_mode == 'graph':
X, I = inputs
else:
X = inputs
attn_coeff = K.dot(X, self.attn_kernel)
attn_coeff = K.squeeze(attn_coeff, -1)
attn_coeff = K.softmax(attn_coeff)
if self.data_mode == 'single':
output = K.dot(attn_coeff[None, ...], X)
elif self.data_mode == 'batch':
output = K.batch_dot(attn_coeff, X)
else:
output = attn_coeff[:, None] * X
output = tf.segment_sum(output, I)
return output
def call(self, inputs):
if self.data_mode == 'graph':
X, I = inputs
else:
X = inputs
inputs_linear = K.dot(X, self.lg_kernel) + self.lg_bias
attn_map = K.dot(X, self.attn_kernel) + self.attn_bias
attn_map = K.sigmoid(attn_map)
masked_inputs = inputs_linear * attn_map
if self.data_mode in {'single', 'batch'}:
output = K.sum(masked_inputs,
axis=-2,
keepdims=self.data_mode == 'single')
else:
output = tf.segment_sum(masked_inputs, I)
return output