def dot_or_tensor_mul(self, features, tensor):
tensor_shape = tf.shape(tensor)
flat_shape = [tensor_shape[0], tensor_shape[1] * tensor_shape[2]]
flattened_tensor = tf.reshape(tensor, flat_shape)
result_tensor = dot_or_lookup(features, flattened_tensor, onehot_input=self.onehot_input)
result_tensor = tf.reshape(result_tensor, [-1, tensor_shape[1], tensor_shape[2]])
return result_tensor
def compute_self_loop_messages(self, vertex_features):
return dot_or_lookup(vertex_features,
self.W_self,
onehot_input=self.onehot_input)
