def enrich(triples: torch.Tensor, n: int, r: int):
cuda = triples.is_cuda
inverses = torch.cat([triples[:, 2:], triples[:, 1:2] + r, triples[:, :1]],
dim=1)
selfloops = torch.cat([
torch.arange(n, dtype=torch.long, device=d(cuda))[:, None],
torch.full((n, 1), fill_value=2 * r),
torch.arange(n, dtype=torch.long, device=d(cuda))[:, None],
],
dim=1)
return torch.cat([triples, inverses, selfloops], dim=0)
def sum_sparse(indices, values, size, row=True):
"""
Sum the rows or columns of a sparse matrix, and redistribute the
results back to the non-sparse row/column entries
:return:
"""
ST = torch.cuda.sparse.FloatTensor if indices.is_cuda else torch.sparse.FloatTensor
assert len(indices.size()) == 2
k, r = indices.size()
if not row:
# transpose the matrix
indices = torch.cat([indices[:, 1:2], indices[:, 0:1]], dim=1)
size = size[1], size[0]
ones = torch.ones((size[1], 1), device=d(indices))
smatrix = ST(indices.t(), values, size=size)
sums = torch.mm(smatrix, ones) # row/column sums
sums = sums[indices[:, 0]]
assert sums.size() == (k, 1)
return sums.view(k)
def adj(triples, num_nodes, num_rels, cuda=False, vertical=True):
"""
Computes a sparse adjacency matrix for the given graph (the adjacency matrices of all
relations are stacked vertically).
:param edges: List representing the triples
:param i2r: list of relations
:param i2n: list of nodes
:return: sparse tensor
"""
r, n = num_rels, num_nodes
size = (r * n, n) if vertical else (n, r * n)
from_indices = []
upto_indices = []
for fr, rel, to in triples:
offset = rel.item() * n
if vertical:
fr = offset + fr.item()
else:
to = offset + to.item()
from_indices.append(fr)
upto_indices.append(to)
indices = torch.tensor([from_indices, upto_indices], dtype=torch.long, device=d(cuda))
assert indices.size(1) == len(triples)
assert indices[0, :].max() < size[0], f'{indices[0, :].max()}, {size}, {r}'
assert indices[1, :].max() < size[1], f'{indices[1, :].max()}, {size}, {r}'
return indices.t(), size