File size: 1,046 Bytes
4f5f309 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | 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))
|