Sampling using low-dimensional random projections (Johnson-Lindenstrauss lemma).
(
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
)
| 691 | return cls(weights, k=k) |
| 692 | |
| 693 | def prepare(self) -> "JLSampler": |
| 694 | D = self.weights.shape[1] # noqa: N806 |
| 695 | self.rand_mat = torch.randn( |
| 696 | (D, self.k), |
| 697 | dtype=self.weights.dtype, |
| 698 | device=self.weights.device, |
| 699 | ) / math.sqrt(self.k) |
| 700 | self.w_p = self.weights @ self.rand_mat # [V, k] |
| 701 | self.w_p = self.w_p.contiguous() |
| 702 | self.prepared = True |
| 703 | self.weights = None # not needed anymore |
| 704 | return self |
| 705 | |
| 706 | @torch.compile(fullgraph=True) |
| 707 | def sample( |
| 708 | self, |
| 709 | hidden_states: torch.Tensor, # [n_hidden_states, D] |
| 710 | temperature: torch.Tensor, # scalar (0-d) |
| 711 | num_samples: int, |
| 712 | seed: int | None = None, # ignored |
nothing calls this directly
no test coverage detected