Computes the logsumexp of each segment in a tensor.
(src, dst_idx, dst_size, extra_dims=None)
| 156 | |
| 157 | |
| 158 | def segment_logsumexp(src, dst_idx, dst_size, extra_dims=None): |
| 159 | """Computes the logsumexp of each segment in a tensor.""" |
| 160 | src_max, _ = scatter_max(src, dst_idx, dim=0, dim_size=dst_size) |
| 161 | if extra_dims is not None: |
| 162 | src_max = torch.amax(src_max, dim=extra_dims, keepdim=True) |
| 163 | src = src - src_max[dst_idx] |
| 164 | out = torch.zeros( |
| 165 | dst_size, |
| 166 | *src.shape[1:], |
| 167 | dtype=src.dtype, |
| 168 | device=src.device, |
| 169 | ).index_add_(0, dst_idx, torch.exp(src)) |
| 170 | if extra_dims is not None: |
| 171 | out = torch.sum(out, dim=extra_dims) |
| 172 | return torch.log(out + 1e-8) + src_max.view(*out.shape) |
| 173 | |
| 174 | |
| 175 | def segment_softmax(src, dst_idx, dst_size, extra_dims=None, floor_value=None): |
no test coverage detected