(seq_k, total_bits)
| 26 | @pytest.mark.parametrize("seq_k", [64, 128, 256, 512]) |
| 27 | @pytest.mark.parametrize("total_bits", [3, 4]) |
| 28 | def test_compress_keys_shapes(seq_k, total_bits): |
| 29 | engine = TurboQuantEngine(head_dim=128, total_bits=total_bits, device="cpu") |
| 30 | K = _make_random_keys(seq_k) |
| 31 | compressed = engine.compress_keys_pytorch(K) |
| 32 | |
| 33 | assert compressed["indices"].shape == (seq_k, 128) |
| 34 | assert compressed["k_mse"].shape == (seq_k, 128) |
| 35 | assert compressed["qjl_signs"].shape == (seq_k, 128) |
| 36 | assert compressed["vec_norms"].shape == (seq_k,) |
| 37 | assert compressed["residual_norms"].shape == (seq_k,) |
| 38 | |
| 39 | |
| 40 | @pytest.mark.parametrize("total_bits", [3, 4]) |
nothing calls this directly
no test coverage detected