def fn(rel_id, num_chunks, head, tail, gpu_id, trace=False): # pos node, project to its relation projection = self.projection_emb(rel_id, gpu_id, trace) projection = projection.reshape(-1, self.entity_dim, self.relation_dim) head = head.reshape(-1, 1, self.entity_dim) head = nd.batch_dot(head, projection).squeeze() head = head.reshape(num_chunks, -1, self.relation_dim) projection = projection.reshape(num_chunks, -1, self.entity_dim, self.relation_dim) tail = tail.reshape(num_chunks, -1, 1, self.entity_dim) num_rels = projection.shape[1] num_nnodes = tail.shape[1] tails = [] for i in range(num_chunks): tail_negs = [] for j in range(num_nnodes): tail_neg = tail[i][j] tail_neg = tail_neg.reshape(1, 1, self.entity_dim) tail_neg = nd.broadcast_axis(tail_neg, axis=0, size=num_rels) tail_neg = nd.batch_dot(tail_neg, projection[i]) tail_neg = tail_neg.squeeze(axis=1) tail_negs.append(tail_neg) tail_negs = nd.stack(*tail_negs, axis=1) tails.append(tail_negs) tail = nd.stack(*tails) return head, tail
def rbf_kernels(self, x: NDArray, y: NDArray): """ Computes exp(-c ||x - y||^2). ||x - y||^2 = x . x + y . y - 2 x . y Compute each term separately. x is are original features, y are features used for similarity """ cross_products = nd.dot(x, y) x_products = nd.sum(sqr(x), axis=1, keepdims=True) x_products = nd.broadcast_axis(x_products, axis=1, size=y.shape[1]) y_products = nd.sum(sqr(y), axis=0, keepdims=True) y_products = nd.broadcast_axis(y_products, axis=0, size=x.shape[0]) sqr_difs = x_products + y_products - 2 * cross_products print(nd.mean(x_products), nd.mean(y_products), nd.mean(cross_products)) print(nd.mean(sqr_difs)) res = nd.exp(-0.05 * sqr_difs) print(res.shape) return res