def test_to_csr(self):
self.data._adj._to_csr()
symmetric = self.data.is_symmetric()
assert symmetric is True
degrees = self.data.degrees()
_degrees = get_degrees(self.data.edge_index)
assert (degrees == _degrees).all()
def forward(self, x, edge_index, edge_attr):
N = x.shape[0]
row, col = edge_index
deg = get_degrees(row, col, N)
deg_inv = deg.pow(-1)
adj = torch.sparse_coo_tensor(edge_index,
deg_inv[row] * edge_attr,
size=(N, N))
identity = torch.sparse_coo_tensor([range(N)] * 2,
torch.ones(N),
size=(N, N)).to(x.device)
laplacian = identity - adj
t0 = identity
t1 = laplacian - self.mu * identity
t1 = t1.mm(t1.to_dense()).to_sparse()
l_x = -0.5 * (t1 - identity)
ivs = [iv(i, self.theta) for i in range(self.steps)]
ivs[1:] = [(-1)**i * 2 * x for i, x in enumerate(ivs[1:])]
ivs = torch.tensor(ivs).to(x.device)
result = [t0, l_x]
for i in range(2, self.steps):
result.append(2 * l_x.mm(result[i - 1].to_dense()).to_sparse().sub(
result[i - 2]))
result = [result[i] * ivs[i] for i in range(self.steps)]
def fn(x, y):
return x.add(y)
res = reduce(fn, result)
return res._indices(), res._values()
def forward(self, x, edge_index, edge_attr):
row, col = edge_index
deg = get_degrees(row, col, x.shape[0])
deg_inv_sqrt = deg.pow(-0.5)
edge_attr_t = deg_inv_sqrt[row] * edge_attr * deg_inv_sqrt[col]
N = x.size(0)
adj = torch.sparse_coo_tensor(edge_index, edge_attr_t, size=(N, N))
theta = self.alpha * (1 - self.alpha)
result = [theta * adj]
for i in range(1, self.steps - 1):
theta = theta * (1 - self.alpha)
adj_ind, adj_val = spspmm(edge_index, edge_attr_t,
result[i - 1]._indices(),
result[i - 1]._values(), N, N, N, True)
result.append(
torch.sparse_coo_tensor(adj_ind, adj_val, size=(N, N)))
identity = torch.sparse_coo_tensor([range(N)] * 2,
torch.ones(N),
size=(N, N)).to(x.device)
result.append(self.alpha * identity)
def fn(x, y):
return x.add(y)
res = reduce(fn, result)
return res._indices(), res._values()
def degrees(self, node_idx=None):
if self.row_ptr is not None:
degs = (self.row_ptr[1:] - self.row_ptr[:-1]).float()
if node_idx is not None:
return degs[node_idx]
return degs
else:
return get_degrees(self.row, self.col, num_nodes=self.num_nodes)
def forward(self, x, edge_index, edge_attr):
device = x.device
N, dim = x.shape
diag_val = self.p(x)
diag_val = torch.sigmoid(diag_val)
self.dropout(diag_val)
row, col = edge_index
deg = get_degrees(row, col, N)
deg_inv = deg.pow(-1)
edge_attr_t = deg_inv[row] * edge_attr
diag_ind = torch.LongTensor([range(N)] * 2).to(device)
_, adj_mat_val = spspmm(edge_index, edge_attr_t, diag_ind,
diag_val.view(-1), N, N, N, True)
return edge_index, adj_mat_val
def forward(self, x, edge_index, edge_attr):
N, dim = x.shape
row, col = edge_index
deg = get_degrees(row, col, N)
deg_inv_sqrt = deg.pow(-0.5)
edge_attr_t = deg_inv_sqrt[row] * edge_attr * deg_inv_sqrt[col]
p_val = F.relu(self.p(x))
q_val = F.relu(self.q(x))
p_val = self.dropout(p_val)
q_val = self.dropout(q_val)
p_adj_mat_val = edge_attr_t * p_val.view(-1)[edge_index[1]]
q_adj_mat_val = edge_attr_t * q_val.view(-1)[edge_index[0]]
return edge_index, p_adj_mat_val + q_adj_mat_val
def degrees(self):
if self.row_ptr is not None:
return self.row_ptr[1:] - self.row_ptr[:-1]
else:
edge_index = torch.stack([self.row, self.col])
return get_degrees(edge_index, num_nodes=self.num_nodes)
def degrees(self):
if self.row_ptr is not None:
return self.row_ptr[1:] - self.row_ptr[:-1]
else:
return get_degrees(self.row, self.col, num_nodes=self.num_nodes)
def forward(self, x, edge_index, edge_attr):
row, col = edge_index
deg = get_degrees(row, col, x.shape[0])
deg_inv = deg.pow(-1)
edge_attr_t = self.t * edge_attr * deg_inv[col] - self.t
return edge_index, edge_attr_t.exp()
