Samples the index of the minimum value in each segment.
(scores, dst_idx, dst_size, randomize: bool = False)
| 147 | |
| 148 | |
| 149 | def segment_argmin(scores, dst_idx, dst_size, randomize: bool = False) -> torch.Tensor: |
| 150 | """Samples the index of the minimum value in each segment.""" |
| 151 | if randomize: |
| 152 | noise = torch.rand_like(scores) |
| 153 | scores = scores - torch.log(-torch.log(noise)) |
| 154 | _, sampled_idx = scatter_min(scores, dst_idx, dim=0, dim_size=dst_size) |
| 155 | return sampled_idx |
| 156 | |
| 157 | |
| 158 | def segment_logsumexp(src, dst_idx, dst_size, extra_dims=None): |
no test coverage detected