Method compute_logits

src/fused_mm_sampling/core.py:711–716

(
        self,
        hidden_states: torch.Tensor,  # [n_hidden_states, D]
    )

        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]

    def rrt(self) -> torch.Tensor:
        """Return R @ Rᵀ, which should be close to the identity matrix."""