Function test_quantize_nearest_centroid

tests/test_codebook.py:86–92

Values close to a centroid should map to that centroid's index.

()

def test_quantize_nearest_centroid():
    """Values close to a centroid should map to that centroid's index."""
    cb = LloydMaxCodebook(128, 2)
    for i, c in enumerate(cb.centroids):
        test_val = c + 1e-6
        idx = cb.quantize(test_val.unsqueeze(0))
        assert idx.item() == i, f"Expected index {i}, got {idx.item()}"