Function optimal_k

src/fused_mm_sampling/core.py:724–727

Source: https://cs.stanford.edu/people/mmahoney/cs369m/Lectures/lecture1.pdf

(n: int, epsilon: float)

        samples = _fast_multinomial(probs, num_samples)
        return samples

    def compute_logits(
        self,
        hidden_states: torch.Tensor,  # [n_hidden_states, D]
    ) -> torch.Tensor:
        h_p = hidden_states @ self.rand_mat  # [n_hidden_states, k]
        return h_p @ self.w_p.T  # [n_hidden_states, V]