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