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
def forward(self, q: QTensor, edge_index: Adj, edge_attr: QTensor, size: Size = None) -> QTensor: assert edge_attr.__class__.__name__ == "QTensor" x = q.clone() # propagate each part, i.e. real part and the three complex parts agg_r = self.propagate(edge_index=edge_index, x=q.r, edge_attr=edge_attr.r, size=size) # [b_num_nodes, in_features] agg_i = self.propagate(edge_index=edge_index, x=q.i, edge_attr=edge_attr.i, size=size) # [b_num_nodes, in_features] agg_j = self.propagate(edge_index=edge_index, x=q.j, edge_attr=edge_attr.j, size=size) # [b_num_nodes, in_features] agg_k = self.propagate(edge_index=edge_index, x=q.k, edge_attr=edge_attr.k, size=size) # [b_num_nodes, in_features] q = QTensor(agg_r, agg_i, agg_j, agg_k) 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