(self, coords, batch_size)
| 121 | return torch.Size(shape) |
| 122 | |
| 123 | def __cal_layout(self, coords, batch_size): |
| 124 | seq_len = torch.bincount(coords[:, 0], minlength=batch_size) |
| 125 | offset = torch.cumsum(seq_len, dim=0) |
| 126 | layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)] |
| 127 | return layout |
| 128 | |
| 129 | @property |
| 130 | def shape(self) -> torch.Size: |