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 | |
| 692 | @torch.compile(fullgraph=True) |
| 693 | def sample( |
| 694 | self, |
| 695 | hidden_states: torch.Tensor, # [n_hidden_states, D] |
| 696 | temperature: torch.Tensor, # scalar (0-d) |
| 697 | num_samples: int, |
| 698 | seed: int | None = None, # ignored |
| 699 | weights: torch.Tensor = None, # ignored |
| 700 | ): |
| 701 | """ |
| 702 | Sampling using low-dimensional random projections (Johnson-Lindenstrauss lemma). |
| 703 | """ |
| 704 | if not self.prepared: |
| 705 | raise ValueError("Sampler not prepared. Call .prepare() first.") |
| 706 | logits_p = self.compute_logits(hidden_states) |
| 707 | probs = (logits_p / temperature).softmax(dim=1) |
| 708 | samples = _fast_multinomial(probs, num_samples) |
| 709 | return samples |
| 710 | |
| 711 | def compute_logits( |
| 712 | self, |
nothing calls this directly
no test coverage detected