Helper function for the diagonal ranges function.
(batch)
| 11 | |
| 12 | |
| 13 | def ranges_slices(batch): |
| 14 | """Helper function for the diagonal ranges function.""" |
| 15 | if 'mps' in str(batch.device): |
| 16 | Ns = batch.to('cpu').bincount() |
| 17 | else: |
| 18 | Ns = batch.bincount() |
| 19 | |
| 20 | indices = Ns.cumsum(0) |
| 21 | ranges = torch.cat((0 * indices[:1], indices)) |
| 22 | ranges = ( |
| 23 | torch.stack((ranges[:-1], ranges[1:])).t().int().contiguous().to(batch.device) |
| 24 | ) |
| 25 | slices = (1 + torch.arange(len(Ns))).int().to(batch.device) |
| 26 | |
| 27 | return ranges, slices |
| 28 | |
| 29 | |
| 30 | def diagonal_ranges(batch_x=None, batch_y=None): |