def panentropy(self,
adj_t: SparseTensor,
dtype: Optional[int] = None) -> SparseTensor:
tmp = SparseTensor.eye(adj_t.size(0),
adj_t.size(1),
has_value=True,
dtype=dtype,
device=adj_t.device())
tmp = tmp.mul_nnz(self.weight[0])
outs = [tmp]
for i in range(1, self.filter_size + 1):
tmp = tmp @ adj_t
tmp = tmp.mul_nnz(self.weight[i])
outs += [tmp]
row = torch.cat([out.storage.row() for out in outs], dim=0)
col = torch.cat([out.storage.col() for out in outs], dim=0)
value = torch.cat([out.storage.value() for out in outs], dim=0)
out = SparseTensor(row=row,
col=col,
value=value,
sparse_sizes=adj_t.sparse_sizes()).coalesce()
return out
def forward(self, x, edge_index, edge_attr=None, batch=None):
if batch is None:
batch = edge_index.new_zeros(x.size(0))
num_node = x.size(0)
k = F.relu(self.lin_2(x))
A = SparseTensor.from_edge_index(edge_index=edge_index,
edge_attr=edge_attr,
sparse_sizes=(num_node, num_node))
I = SparseTensor.eye(num_node, device=self.args.device)
A_wave = fill_diag(A, 1)
s = A_wave @ k
score = s.squeeze()
perm = topk(score, self.ratio, batch)
A = self.norm(A)
K_neighbor = A * k.T
x_neighbor = K_neighbor @ x
# ----modified
deg = sum(A, dim=1)
deg_inv = deg.pow_(-1)
deg_inv.masked_fill_(deg_inv == float('inf'), 0.)
x_neighbor = x_neighbor * deg_inv.view(1, -1).T
# ----
x_self = x * k
x = x_neighbor * (
1 - self.args.combine_ratio) + x_self * self.args.combine_ratio
x = x[perm]
batch = batch[perm]
edge_index, edge_attr = filter_adj(edge_index,
edge_attr,
perm,
num_nodes=s.size(0))
return x, edge_index, edge_attr, batch, perm
