Sample from a categorical distribution using the exponential race method. Avoids torch.multinomial's 10-kernel validation overhead (~2/3 of its runtime). For each row, draws exponential noise, computes probs / noise, and takes argmax. See https://github.com/pytorch/pytorch/issues/177127
(probs: torch.Tensor, num_samples: int)
| 62 | |
| 63 | |
| 64 | def _fast_multinomial(probs: torch.Tensor, num_samples: int) -> torch.Tensor: |
| 65 | """Sample from a categorical distribution using the exponential race method. |
| 66 | |
| 67 | Avoids torch.multinomial's 10-kernel validation overhead (~2/3 of its runtime). |
| 68 | For each row, draws exponential noise, computes probs / noise, and takes argmax. |
| 69 | See https://github.com/pytorch/pytorch/issues/177127 |
| 70 | """ |
| 71 | H, V = probs.shape # noqa: N806 |
| 72 | q = torch.empty(num_samples, H, V, device=probs.device, dtype=probs.dtype) |
| 73 | q.exponential_() |
| 74 | return probs.unsqueeze(0).div(q).argmax(dim=-1).T # [H, num_samples] |
| 75 | |
| 76 | |
| 77 | def _allgather_logits( |