Method rrt

src/fused_mm_sampling/core.py:718–721

Return R @ Rᵀ, which should be close to the identity matrix.

(self)

        Sampling using low-dimensional random projections (Johnson-Lindenstrauss lemma).
        """
        if not self.prepared:
            raise ValueError("Sampler not prepared. Call .prepare() first.")
        logits_p = self.compute_logits(hidden_states)
        probs = (logits_p / temperature).softmax(dim=1)
        samples = _fast_multinomial(probs, num_samples)
        return samples