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