Function print_once

src/fused_mm_sampling/core.py:109–110

(msg: str)

    out = torch.zeros_like(logits)
    return out.scatter_(dim=-1, index=topk_idx, src=probs)


@lru_cache
def print_once(msg: str):