()
| 21 | |
| 22 | |
| 23 | def test_jl_sampling_aproximate_correctness(): |
| 24 | folder = Path(__file__).parent / "qwen3-0.6b" |
| 25 | assert folder.exists(), f"{folder} does not exist. Run generate_inputs.py first" |
| 26 | weights = torch.load(folder / "weights.pt", map_location=device) |
| 27 | hidden_states = torch.load(folder / "hidden_states.pt", map_location=device) |
| 28 | expected_logits = hidden_states @ weights.T |
| 29 | expected_probs = expected_logits.softmax(dim=1) |
| 30 | |
| 31 | jl_sampler = JLSampler.from_weights(weights, epsilon=0.2).prepare() |
| 32 | actual_logits = jl_sampler.compute_logits(hidden_states) |
| 33 | actual_probs = actual_logits.softmax(dim=1) |
| 34 | assert torch.allclose(actual_probs, expected_probs, atol=0.2) |
| 35 | |
| 36 | print(f"{weights.shape=}") |
| 37 | print(f"{hidden_states.shape=}") |
| 38 | |
| 39 | |
| 40 | @pytest.mark.parametrize( |
nothing calls this directly
no test coverage detected