| from __future__ import annotations | |
| import numpy as np | |
| def gumbel(rng: np.random.Generator, shape: tuple[int, ...]) -> np.ndarray: | |
| """ | |
| Produce float32 Gumbel noise: -log(-log(U)) with U in [1e-10, 1 - 1e-7]. | |
| """ | |
| u = rng.uniform(low=1e-10, high=1.0 - 1e-7, size=shape).astype(np.float32) | |
| return (-np.log(-np.log(u))).astype(np.float32) | |
| def topk_gumbel_argmax( | |
| logits: np.ndarray, | |
| temp: float = 0.7, | |
| top_k: int = 64, | |
| noise: np.ndarray | None = None, | |
| ) -> int: | |
| """ | |
| Numpy reference of the graph's rule: | |
| argmax(where(rank < top_k, logits / temp + noise, -inf)) | |
| """ | |
| logits = np.asarray(logits, dtype=np.float32) | |
| if noise is None: | |
| noise = np.zeros_like(logits) | |
| else: | |
| noise = np.asarray(noise, dtype=np.float32) | |
| k = min(top_k, logits.shape[-1]) | |
| topk_idx = np.argpartition(logits, -k)[-k:] | |
| masked_perturbed = np.full_like(logits, -np.inf) | |
| masked_perturbed[topk_idx] = logits[topk_idx] / temp + noise[topk_idx] | |
| return int(np.argmax(masked_perturbed)) | |