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