def forward(
self, y: Tensor, edge_index: Adj, mask: Optional[Tensor] = None,
edge_weight: OptTensor = None,
post_step: Callable = lambda y: y.clamp_(0., 1.)
) -> Tensor:
""""""
if y.dtype == torch.long:
y = F.one_hot(y.view(-1)).to(torch.float)
out = y
if mask is not None:
out = torch.zeros_like(y)
out[mask] = y[mask]
if isinstance(edge_index, SparseTensor) and not edge_index.has_value():
edge_index = gcn_norm(edge_index, add_self_loops=False)
elif isinstance(edge_index, Tensor) and edge_weight is None:
edge_index, edge_weight = gcn_norm(edge_index, num_nodes=y.size(0),
add_self_loops=False)
res = (1 - self.alpha) * out
for _ in range(self.num_layers):
# propagate_type: (y: Tensor, edge_weight: OptTensor)
out = self.propagate(edge_index, x=out, edge_weight=edge_weight,
size=None)
out.mul_(self.alpha).add_(res)
out = post_step(out)
return out
def forward(self,
y: Tensor,
edge_index: Adj,
mask: Tensor,
edge_weight: OptTensor = None) -> Tensor:
""""""
if mask.dtype == torch.long:
idx = mask
mask = mask.new_zeros(y.size(0), dtype=torch.bool)
mask[idx] = True
y = F.one_hot(y.view(-1)).to(torch.float)
y[~mask] = 0.
if isinstance(edge_index, SparseTensor) and not edge_index.has_value():
edge_index = gcn_norm(edge_index, add_self_loops=False)
elif isinstance(edge_index, Tensor) and edge_weight is None:
edge_index, edge_weight = gcn_norm(edge_index,
num_nodes=y.size(0),
add_self_loops=False)
res = (1 - self.alpha) * y
for _ in range(self.num_layers):
# propagate_type: (y: Tensor, edge_weight: OptTensor)
y = self.propagate(edge_index,
y=y,
edge_weight=edge_weight,
size=None)
y.mul_(self.alpha).add_(res).clamp_(0, 1)
return y
def forward(self,
x: Tensor,
x_0: Tensor,
edge_index: Adj,
edge_weight: OptTensor = None,
return_attention_weights=None):
# type: (Tensor, Tensor, Tensor, OptTensor, NoneType) -> Tensor # noqa
# type: (Tensor, Tensor, SparseTensor, OptTensor, NoneType) -> Tensor # noqa
# type: (Tensor, Tensor, Tensor, OptTensor, bool) -> Tuple[Tensor, Tuple[Tensor, Tensor]] # noqa
# type: (Tensor, Tensor, SparseTensor, OptTensor, bool) -> Tuple[Tensor, SparseTensor] # noqa
r"""
Args:
return_attention_weights (bool, optional): If set to :obj:`True`,
will additionally return the tuple
:obj:`(edge_index, attention_weights)`, holding the computed
attention weights for each edge. (default: :obj:`None`)
"""
if self.normalize:
if isinstance(edge_index, Tensor):
assert edge_weight is None
cache = self._cached_edge_index
if cache is None:
edge_index, edge_weight = gcn_norm( # yapf: disable
edge_index,
None,
x.size(self.node_dim),
False,
self.add_self_loops,
dtype=x.dtype)
if self.cached:
self._cached_edge_index = (edge_index, edge_weight)
else:
edge_index, edge_weight = cache[0], cache[1]
elif isinstance(edge_index, SparseTensor):
assert not edge_index.has_value()
cache = self._cached_adj_t
if cache is None:
edge_index = gcn_norm( # yapf: disable
edge_index,
None,
x.size(self.node_dim),
False,
self.add_self_loops,
dtype=x.dtype)
if self.cached:
self._cached_adj_t = edge_index
else:
edge_index = cache
else:
if isinstance(edge_index, Tensor):
assert edge_weight is not None
elif isinstance(edge_index, SparseTensor):
assert edge_index.has_value()
alpha_l = self.att_l(x)
alpha_r = self.att_r(x)
# propagate_type: (x: Tensor, alpha: PairTensor, edge_weight: OptTensor) # noqa
out = self.propagate(edge_index,
x=x,
alpha=(alpha_l, alpha_r),
edge_weight=edge_weight,
size=None)
alpha = self._alpha
self._alpha = None
if self.eps != 0.0:
out += self.eps * x_0
if isinstance(return_attention_weights, bool):
assert alpha is not None
if isinstance(edge_index, Tensor):
return out, (edge_index, alpha)
elif isinstance(edge_index, SparseTensor):
return out, edge_index.set_value(alpha, layout='coo')
else:
return out