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
