Module sampling 

Gumbel-Softmax sampling and relaxed top-k. Gumbel-Softmax sampling and relaxed top-k.

Delegates to kuji for correct implementations of the Gumbel-Softmax trick and the iterated masked-softmax relaxed top-k (Kool et al., 2019).

Requires the gumbel feature (enables rand + kuji dependencies).

Functions

gumbel_attention_mask

Generate Gumbel-based attention mask for RAG reranking.

gumbel_softmax

Gumbel-Softmax: differentiable sampling from categorical distribution.

relaxed_topk_gumbel

Relaxed Top-k using iterated Gumbel-Softmax (Kool et al., 2019).