def forward(self,
q: QTensor,
edge_index: Adj,
edge_attr: QTensor,
size: Size = None) -> QTensor:
assert edge_attr.__class__.__name__ == "QTensor"
x = q.clone()
# "cast" QTensor back to torch.Tensor
q = q.stack(dim=1) # (batch_num_nodes, 4, feature_dim)
q = q.reshape(q.size(0), -1) # (batch_num_nodes, 4*feature_dim)
edge_attr = edge_attr.stack(dim=1)
edge_attr = edge_attr.reshape(edge_attr.size(0), -1)
# propagate
agg = self.propagate(edge_index=edge_index,
x=q,
edge_attr=edge_attr,
size=size)
agg = agg.reshape(agg.size(0), 4, -1).permute(1, 0, 2)
agg = QTensor(*agg)
if self.add_self_loops:
x += agg
# transform aggregated node embeddings
q = self.transform(x)
return q
def forward(self,
q: QTensor,
edge_index: Adj,
edge_attr: QTensor,
size: Size = None) -> QTensor:
assert edge_attr.__class__.__name__ == "QTensor"
x = q.clone()
# "cast" QTensor back to torch.Tensor
q = q.stack(dim=1) # (batch_num_nodes, 4, feature_dim)
q = q.reshape(q.size(0), -1) # (batch_num_nodes, 4*feature_dim)
edge_attr = edge_attr.stack(dim=1)
edge_attr = edge_attr.reshape(edge_attr.size(0), -1)
# propagate
agg = self.propagate(edge_index=edge_index,
x=q,
edge_attr=edge_attr,
size=size)
agg = agg.reshape(agg.size(0), 4, -1).permute(1, 0, 2)
q = QTensor(*agg)
if self.same_dim: # aggregate messages -> linearly transform -> add self-loops.
q = self.transform(q)
if self.add_self_loops:
q += x
else:
if self.add_self_loops: # aggregate messages -> add self-loops -> linearly transform.
q += x
q = self.transform(q)
return q