| 242 | |
| 243 | |
| 244 | class Scaling_SparseEdge_Multi(SparseEdge_Multi): |
| 245 | def __init__(self, num_from: int, max_deg: int): |
| 246 | self.out_deg = AutoScalingTensor((num_from,), grow_on=0, dtype=torch.long, init_val=0) |
| 247 | self.edges = AutoScalingTensor((num_from, max_deg), grow_on=0, dtype=torch.long, init_val=-1) |
| 248 | self.max_deg = max_deg |
| 249 | |
| 250 | |
| 251 | def push(self, edge_multi: SparseEdge_Multi): |
| 252 | self.out_deg.push(edge_multi.out_deg) |
| 253 | self.edges .push(edge_multi.edges) |
| 254 | |
| 255 | @classmethod |
| 256 | def deserialize(cls, prefix: str, value: dict[str, np.ndarray]) -> Self: |
| 257 | tensor_edge = super().deserialize(prefix, value) |
| 258 | |
| 259 | # Convert from torch.Tensor to AutoScalingTensor |
| 260 | tensor_edge.out_deg = AutoScalingTensor(None, grow_on=0, init_tensor=tensor_edge.out_deg) |
| 261 | tensor_edge.edges = AutoScalingTensor(None, grow_on=0, init_tensor=tensor_edge.edges) |
| 262 | return tensor_edge |
| 263 | |
| 264 | |
| 265 | class Scaling_DenseEdge_Multi(DenseEdge_Multi): |