def message(self, x_j: Tensor, alpha_j: Tensor, alpha_i: OptTensor,
index: Tensor, ptr: OptTensor, edge_weight: OptTensor,
size_i: Optional[int]) -> Tensor:
alpha = alpha_j if alpha_i is None else alpha_j + alpha_i
alpha = F.leaky_relu(alpha, self.negative_slope)
alpha = softmax(alpha, index, ptr, size_i)
alpha = self.att_norm(alpha * edge_weight.unsqueeze(-1), index)
self._alpha = alpha
alpha = F.dropout(alpha, p=self.dropout, training=self.training)
return x_j * alpha.unsqueeze(-1)
def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
assert edge_weight is not None
return x_j * edge_weight.unsqueeze(1)
