def _call_single(self, x, a): # Reshape kernels for efficient message-passing kernel = tf.reshape(self.kernel, (-1, self.attn_heads * self.channels)) attn_kernel_self = ops.transpose(self.attn_kernel_self, (2, 1, 0)) attn_kernel_neighs = ops.transpose(self.attn_kernel_neighs, (2, 1, 0)) # Prepare message-passing indices = a.indices N = tf.shape(x, out_type=indices.dtype)[-2] indices = ops.add_self_loops_indices(indices, N) targets, sources = indices[:, 1], indices[:, 0] # Update node features x = K.dot(x, kernel) x = tf.reshape(x, (-1, self.attn_heads, self.channels)) # Compute attention attn_for_self = tf.reduce_sum(x * attn_kernel_self, -1) attn_for_self = tf.gather(attn_for_self, targets) attn_for_neighs = tf.reduce_sum(x * attn_kernel_neighs, -1) attn_for_neighs = tf.gather(attn_for_neighs, sources) attn_coef = attn_for_self + attn_for_neighs attn_coef = tf.nn.leaky_relu(attn_coef, alpha=0.2) attn_coef = ops.unsorted_segment_softmax(attn_coef, targets, N) attn_coef = self.dropout(attn_coef) attn_coef = attn_coef[..., None] # Update representation output = attn_coef * tf.gather(x, sources) output = tf.math.unsorted_segment_sum(output, targets, N) return output, attn_coef
def _call_single(self, X, A): # Reshape kernels for efficient message-passing kernel = tf.reshape(self.kernel, (-1, self.attn_heads * self.channels)) attn_kernel_self = ops.transpose(self.attn_kernel_self, (2, 1, 0)) attn_kernel_neighs = ops.transpose(self.attn_kernel_neighs, (2, 1, 0)) # Prepare message-passing indices = A.indices N = tf.shape(X, out_type=indices.dtype)[0] indices = ops.sparse_add_self_loops(indices, N) targets, sources = indices[:, -2], indices[:, -1] # Update node features X = ops.dot(X, kernel) X = tf.reshape(X, (-1, self.attn_heads, self.channels)) # Compute attention attn_for_self = tf.reduce_sum(X * attn_kernel_self, -1) attn_for_self = tf.gather(attn_for_self, targets) attn_for_neighs = tf.reduce_sum(X * attn_kernel_neighs, -1) attn_for_neighs = tf.gather(attn_for_neighs, sources) attn_coef = attn_for_self + attn_for_neighs attn_coef = tf.nn.leaky_relu(attn_coef, alpha=0.2) attn_coef = ops.unsorted_segment_softmax(attn_coef, targets, N) attn_coef = self.dropout(attn_coef) attn_coef = attn_coef[..., None] # Update representation output = attn_coef * tf.gather(X, sources) output = ops.scatter_sum(targets, output, N) return output, attn_coef