|
25 | 25 | """ |
26 | 26 |
|
27 | 27 | import os |
28 | | -import unittest |
29 | 28 | from typing import Callable, Dict, List, Optional, Tuple |
30 | 29 |
|
31 | 30 | import torch |
@@ -7703,95 +7702,43 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: |
7703 | 7702 | ) |
7704 | 7703 |
|
7705 | 7704 |
|
7706 | | -def _ref_gumbel_max(logits: torch.Tensor, temperature: float, seed: int): |
7707 | | - """Independent Gumbel-max reference using the same torch RNG as the op.""" |
7708 | | - gen = torch.Generator().manual_seed(seed) |
7709 | | - u = torch.rand(logits.shape, generator=gen) |
7710 | | - gumbel = -torch.log(-torch.log(u)) |
7711 | | - return torch.argmax(logits / temperature + gumbel, dim=-1) |
7712 | | - |
7713 | | - |
7714 | | -def _tv_distance(p: torch.Tensor, q: torch.Tensor) -> float: |
7715 | | - """Total-variation distance between two discrete distributions.""" |
7716 | | - return 0.5 * torch.abs(p - q).sum().item() |
7717 | | - |
7718 | | - |
7719 | | -def _sample(logits, temperature, seed: Optional[int], top_p: float = 1.0): |
7720 | | - t = torch.tensor(float(temperature)) |
7721 | | - s = None if seed is None else torch.tensor(int(seed), dtype=torch.int64) |
7722 | | - p = torch.tensor(float(top_p)) # 1.0 = off |
7723 | | - return torch.ops.mlx.sample(logits, t, p, s) |
7724 | | - |
7725 | | - |
7726 | | -class TestSampleOp(unittest.TestCase): |
7727 | | - """Eager reference behavior of mlx::sample (no export / no runtime).""" |
7728 | | - |
7729 | | - def test_greedy_parity_small_temperature(self): |
7730 | | - # Small temperature -> Gumbel-max collapses to argmax(logits). |
7731 | | - torch.manual_seed(0) |
7732 | | - logits = torch.randn(8, 1024) |
7733 | | - token = _sample(logits, 1e-4, seed=0) |
7734 | | - self.assertTrue(torch.equal(token, torch.argmax(logits, dim=-1))) |
7735 | | - |
7736 | | - def test_greedy_temperature_zero(self): |
7737 | | - # temperature == 0 is exact greedy: argmax(logits), no RNG, no division. |
7738 | | - torch.manual_seed(0) |
7739 | | - logits = torch.randn(8, 1024) |
7740 | | - token = _sample(logits, 0.0, seed=0) |
7741 | | - self.assertTrue(torch.equal(token, torch.argmax(logits, dim=-1))) |
7742 | | - |
7743 | | - def test_matches_independent_gumbel_reference(self): |
7744 | | - # Same seed -> bit-identical token vs an independent Gumbel-max impl. |
7745 | | - torch.manual_seed(1) |
7746 | | - logits = torch.randn(8, 512) |
7747 | | - for seed in (0, 1, 7, 42): |
7748 | | - got = _sample(logits, 0.8, seed=seed) |
7749 | | - expected = _ref_gumbel_max(logits, 0.8, seed) |
7750 | | - self.assertTrue(torch.equal(got, expected), f"mismatch at seed={seed}") |
7751 | | - |
7752 | | - def test_distribution_matches_softmax(self): |
7753 | | - # Empirical token frequencies match softmax(logits / T). |
7754 | | - vocab = 5 |
7755 | | - temperature = 1.0 |
7756 | | - torch.manual_seed(0) |
7757 | | - base = torch.randn(vocab) |
7758 | | - n = 20000 |
7759 | | - tokens = _sample(base.expand(n, vocab), temperature, seed=0) |
7760 | | - |
7761 | | - empirical = torch.bincount(tokens, minlength=vocab).float() / n |
7762 | | - target = torch.softmax(base / temperature, dim=-1) |
7763 | | - tv = _tv_distance(empirical, target) |
7764 | | - self.assertLess(tv, 0.05, f"TV distance {tv:.4f} too large") |
7765 | | - |
7766 | | - def test_determinism_seeded(self): |
7767 | | - # Same seed -> identical draws; different seed -> different draws. |
7768 | | - torch.manual_seed(0) |
7769 | | - logits = torch.randn(256, 64) |
7770 | | - a = _sample(logits, 1.0, seed=123) |
7771 | | - b = _sample(logits, 1.0, seed=123) |
7772 | | - c = _sample(logits, 1.0, seed=124) |
7773 | | - self.assertTrue(torch.equal(a, b)) |
7774 | | - self.assertFalse(torch.equal(a, c)) |
7775 | | - |
7776 | | - def test_unseeded_varies_across_calls(self): |
7777 | | - # seed=None uses the global RNG -> draws vary, tokens stay in range. |
7778 | | - torch.manual_seed(0) |
7779 | | - logits = torch.randn(256, 64) |
7780 | | - a = _sample(logits, 1.0, seed=None) |
7781 | | - b = _sample(logits, 1.0, seed=None) |
7782 | | - self.assertFalse(torch.equal(a, b)) |
7783 | | - self.assertGreaterEqual(int(a.min()), 0) |
7784 | | - self.assertLess(int(a.max()), 64) |
7785 | | - |
7786 | | - def test_top_p_restricts_to_nucleus(self): |
7787 | | - # probs [0.5, 0.3, 0.15, 0.05]; top_p=0.9 keeps {0,1,2}, drops index 3. |
7788 | | - base = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])) |
7789 | | - tokens = _sample(base.expand(5000, 4), 1.0, seed=0, top_p=0.9) |
7790 | | - self.assertTrue((tokens != 3).all()) # tail token never drawn |
7791 | | - self.assertEqual(set(tokens.tolist()), {0, 1, 2}) # nucleus covered |
7792 | | - |
7793 | | - def test_top_p_one_keeps_all(self): |
7794 | | - # top_p=1.0 -> no filtering; the tail token (index 3) is reachable. |
7795 | | - base = torch.log(torch.tensor([0.5, 0.3, 0.15, 0.05])) |
7796 | | - tokens = _sample(base.expand(20000, 4), 1.0, seed=0, top_p=1.0) |
7797 | | - self.assertTrue((tokens == 3).any()) |
| 7705 | +@register_test |
| 7706 | +class SampleGreedyTest(OpTestCase): |
| 7707 | + """Greedy argmax(logits) is bit-exact host/device, so verify the token with the |
| 7708 | + normal compare. Covers temperature=0, tiny temperature, and bf16 logits.""" |
| 7709 | + |
| 7710 | + name = "sample_greedy" |
| 7711 | + |
| 7712 | + def __init__(self, temperature: float = 0.0, dtype: torch.dtype = torch.float32): |
| 7713 | + self.temperature = temperature |
| 7714 | + self.dtype = dtype |
| 7715 | + if dtype == torch.bfloat16: |
| 7716 | + self.name = "sample_greedy_bf16" |
| 7717 | + elif temperature < 0: |
| 7718 | + self.name = "sample_greedy_neg" |
| 7719 | + elif temperature == 0.0: |
| 7720 | + self.name = "sample_greedy" |
| 7721 | + else: |
| 7722 | + self.name = "sample_greedy_eps" |
| 7723 | + |
| 7724 | + @classmethod |
| 7725 | + def get_test_configs(cls) -> List["SampleGreedyTest"]: |
| 7726 | + return [ |
| 7727 | + cls(temperature=0.0), |
| 7728 | + cls(temperature=1e-4), |
| 7729 | + cls(temperature=-1.0), # negative -> greedy on both paths (consistent) |
| 7730 | + cls(temperature=1e-4, dtype=torch.bfloat16), |
| 7731 | + ] |
| 7732 | + |
| 7733 | + def create_model(self) -> nn.Module: |
| 7734 | + return SeededSampleModel() |
| 7735 | + |
| 7736 | + def create_inputs(self) -> Tuple[torch.Tensor, ...]: |
| 7737 | + logits = torch.randn(1, 4, 1024, dtype=self.dtype) |
| 7738 | + if self.dtype == torch.bfloat16: |
| 7739 | + logits[0, -1, 512] = 50.0 # dominant -> unambiguous bf16 argmax |
| 7740 | + return ( |
| 7741 | + logits, |
| 7742 | + torch.tensor(self.temperature), |
| 7743 | + torch.tensor(0, dtype=torch.int64), |
| 7744 | + ) |
0 commit comments