def edge_sum(batch: GraphBatch, copy=True):
if copy:
batch = batch.copy()
batch.e = torch.cat(
[batch.e, batch.e.sum(axis=1, keepdim=True)], axis=1)
return batch
def global_sum(batch: GraphBatch, copy=True):
if copy:
batch = batch.copy()
batch.g = torch.cat(
[batch.g, batch.g.sum(axis=1, keepdim=True)], axis=1)
return batch
def node_sum(batch: GraphBatch, copy=True):
if copy:
batch = batch.copy()
batch.x = torch.cat(
[batch.x, batch.x.sum(axis=1, keepdim=True)], axis=1)
return batch
