Class JLSampler

src/fused_mm_sampling/core.py:664–721

    def sample(self, **kwargs) -> torch.Tensor:
        raise NotImplementedError()


@dataclass
class SimpleSampler(Sampler):
    fn: Callable[..., torch.Tensor]

    def prepare(self) -> "SimpleSampler":
        return self

    def sample(self, **kwargs) -> torch.Tensor:
        return self.fn(**kwargs)


@dataclass
class JLSampler(Sampler):
    weights: torch.Tensor  # [V, D]
    k: int
    prepared: bool = False

    @classmethod
    def from_weights(
        cls,
        weights: torch.Tensor,  # [V, D]
        epsilon: float = 0.2,
    ) -> "JLSampler":
        k = optimal_k(n=weights.shape[0], epsilon=epsilon)
        print(f"JLSampler optimal k={k}")
        return cls(weights, k=k)

    def prepare(self) -> "JLSampler":
        D = self.weights.shape[1]  # noqa: N806
        self.rand_mat = torch.randn(
            (D, self.k),
            dtype=self.weights.dtype,
            device=self.weights.device,
        ) / math.sqrt(self.k)
        self.w_p = self.weights @ self.rand_mat  # [V, k]
        self.w_p = self.w_p.contiguous()
        self.prepared = True
        self.weights = None  # not needed anymore
        return self

    @torch.compile(fullgraph=True)
    def sample(
        self,
        hidden_states: torch.Tensor,  # [n_hidden_states, D]
        temperature: torch.Tensor,  # scalar (0-d)
        num_samples: int,
        seed: int | None = None,  # ignored
        weights: torch.Tensor = None,  # ignored
    ):
        """
        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)