Encodes the block-diagonal structure associated to a batch vector.
(batch_x=None, batch_y=None)
| 28 | |
| 29 | |
| 30 | def diagonal_ranges(batch_x=None, batch_y=None): |
| 31 | """Encodes the block-diagonal structure associated to a batch vector.""" |
| 32 | if batch_x is None and batch_y is None: |
| 33 | return None # No batch processing |
| 34 | elif batch_y is None: |
| 35 | batch_y = batch_x # "symmetric" case |
| 36 | |
| 37 | ranges_x, slices_x = ranges_slices(batch_x) |
| 38 | ranges_y, slices_y = ranges_slices(batch_y) |
| 39 | |
| 40 | return ranges_x, slices_x, ranges_y, ranges_y, slices_y, ranges_x |
| 41 | |
| 42 | |
| 43 | def knn_atoms(x, y, k, batch_x=None, batch_y=None): |
no test coverage detected