make sure to assert out invalid config

Author: lucidrains

pyproject.toml

@@ -41,15 +41,17 @@ build-backend = "hatchling.build"
 
 [tool.rye]
 managed = true
+
 dev-dependencies = [
     "ruff>=0.4.2",
     "pytest>=8.2.0",
     "pytest-cov>=5.0.0",
+    "x-transformers>=2.16.1",
 ]
 
 [tool.pytest.ini_options]
 pythonpath = [
-  "."
+    "."
 ]
 
 [tool.hatch.metadata]

tests/test_lfq.py

@@ -1,77 +1,77 @@
 import torch
 import pytest
-from vector_quantize_pytorch import LFQ
 import math
-"""
-testing_strategy:
-subdivisions: using masks, using frac_per_sample_entropy < 1
-"""
+from vector_quantize_pytorch import LFQ
+
+# helpers
 
-torch.manual_seed(0)
+def exists(val):
+    return val is not None
+
+# tests
 
 @pytest.mark.parametrize('frac_per_sample_entropy', (1., 0.5))
-@pytest.mark.parametrize('mask', (torch.tensor([False, False]),
-                                  torch.tensor([True, False]),
-                                  torch.tensor([True, True])))
+@pytest.mark.parametrize('mask', (
+    torch.tensor([False, False]),
+    torch.tensor([True, False]),
+    torch.tensor([True, True])
+))
 def test_masked_lfq(
     frac_per_sample_entropy,
     mask
 ):
-    # you can specify either dim or codebook_size
-    # if both specified, will be validated against each other
-
     quantizer = LFQ(
-        codebook_size = 65536,      # codebook size, must be a power of 2
-        dim = 16,                   # this is the input feature dimension, defaults to log2(codebook_size) if not defined
-        entropy_loss_weight = 0.1,  # how much weight to place on entropy loss
-        diversity_gamma = 1.,       # within entropy loss, how much weight to give to diversity 
+        codebook_size = 65536,
+        dim = 16,
+        entropy_loss_weight = 0.1,
+        diversity_gamma = 1.,
         frac_per_sample_entropy = frac_per_sample_entropy
     )
 
     image_feats = torch.randn(2, 16, 32, 32)
 
-    ret, loss_breakdown = quantizer(image_feats, inv_temperature=100., return_loss_breakdown=True, mask=mask)  # you may want to experiment with temperature
+    ret, _ = quantizer(image_feats, inv_temperature = 100., return_loss_breakdown = True, mask = mask)
 
     quantized, indices, _ = ret
     assert (quantized == quantizer.indices_to_codes(indices)).all()
 
 @pytest.mark.parametrize('frac_per_sample_entropy', (0.1,))
 @pytest.mark.parametrize('iters', (10,))
 @pytest.mark.parametrize('mask', (None, torch.tensor([True, False])))
-def test_lfq_bruteforce_frac_per_sample_entropy(frac_per_sample_entropy, iters, mask):
+def test_lfq_bruteforce_frac_per_sample_entropy(
+    frac_per_sample_entropy,
+    iters,
+    mask
+):
     image_feats = torch.randn(2, 16, 32, 32)
 
     full_per_sample_entropy_quantizer = LFQ(
-        codebook_size = 65536,      # codebook size, must be a power of 2
-        dim = 16,                   # this is the input feature dimension, defaults to log2(codebook_size) if not defined
-        entropy_loss_weight = 0.1,  # how much weight to place on entropy loss
-        diversity_gamma = 1.,       # within entropy loss, how much weight to give to diversity 
+        codebook_size = 65536,
+        dim = 16,
+        entropy_loss_weight = 0.1,
+        diversity_gamma = 1.,
         frac_per_sample_entropy = 1
     )
 
     partial_per_sample_entropy_quantizer = LFQ(
-        codebook_size = 65536,      # codebook size, must be a power of 2
-        dim = 16,                   # this is the input feature dimension, defaults to log2(codebook_size) if not defined
-        entropy_loss_weight = 0.1,  # how much weight to place on entropy loss
-        diversity_gamma = 1.,       # within entropy loss, how much weight to give to diversity 
+        codebook_size = 65536,
+        dim = 16,
+        entropy_loss_weight = 0.1,
+        diversity_gamma = 1.,
         frac_per_sample_entropy = frac_per_sample_entropy
     )
 
-    ret, loss_breakdown = full_per_sample_entropy_quantizer(
-        image_feats, inv_temperature=100., return_loss_breakdown=True, mask=mask)
+    ret, loss_breakdown = full_per_sample_entropy_quantizer(image_feats, inv_temperature = 100., return_loss_breakdown = True, mask = mask)
     true_per_sample_entropy = loss_breakdown.per_sample_entropy
 
     per_sample_losses = torch.zeros(iters)
-    for iter in range(iters):
-        ret, loss_breakdown = partial_per_sample_entropy_quantizer(
-            image_feats, inv_temperature=100., return_loss_breakdown=True, mask=mask)  # you may want to experiment with temperature
+
+    for i in range(iters):
+        ret, loss_breakdown = partial_per_sample_entropy_quantizer(image_feats, inv_temperature = 100., return_loss_breakdown = True, mask = mask)
 
         quantized, indices, _ = ret
         assert (quantized == partial_per_sample_entropy_quantizer.indices_to_codes(indices)).all()
-        per_sample_losses[iter] = loss_breakdown.per_sample_entropy
-    # 95% confidence interval
-    assert abs(per_sample_losses.mean() - true_per_sample_entropy) \
-        < (1.96*(per_sample_losses.std() / math.sqrt(iters)))
+        per_sample_losses[i] = loss_breakdown.per_sample_entropy
 
-    print("difference: ", abs(per_sample_losses.mean() - true_per_sample_entropy))
-    print("std error:", (1.96*(per_sample_losses.std() / math.sqrt(iters))))
+    # 95% confidence interval
+    assert abs(per_sample_losses.mean() - true_per_sample_entropy) < (1.96 * (per_sample_losses.std() / math.sqrt(iters)))

tests/test_readme.py

@@ -506,3 +506,31 @@ def test_vq_3d():
     assert x.shape == quantized.shape
     assert indices.shape == (1, 16, 16, 16)
     assert torch.allclose(quantized, quantizer.get_output_from_indices(indices))
+
+def test_fvq():
+    from vector_quantize_pytorch import VectorQuantize
+    from x_transformers import ContinuousTransformerWrapper, Encoder
+
+    vq_bridge = ContinuousTransformerWrapper(
+        attn_layers = Encoder(
+            dim = 256,
+            depth = 1,
+            heads = 4,
+            pre_norm_has_final_norm = False
+        )
+    )
+
+    vq = VectorQuantize(
+        dim = 256,
+        codebook_size = 512,
+        vq_bridge = vq_bridge,
+        learnable_codebook = True,
+        ema_update = False,
+        in_place_codebook_optimizer = lambda params: torch.optim.SGD(params, lr = 1e-3)
+    )
+
+    x = torch.randn(1, 1024, 256)
+    quantized, indices, commit_loss = vq(x)
+
+    assert quantized.shape == x.shape
+    assert indices.shape == (1, 1024)

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -813,8 +813,8 @@ def __init__(
 
         # defaults
 
-        ema_update = default(ema_update, not directional_reparam)
-        learnable_codebook = default(learnable_codebook, directional_reparam)
+        ema_update = default(ema_update, not directional_reparam and not exists(vq_bridge))
+        learnable_codebook = default(learnable_codebook, directional_reparam or exists(vq_bridge))
         rotation_trick = default(rotation_trick, not directional_reparam and dim > 1) # only use rotation trick if feature dimension greater than 1
 
         # basic variables
@@ -865,6 +865,8 @@ def __init__(
 
         assert not (straight_through and learnable_codebook), 'gumbel straight through not allowed when learning the codebook'
         assert not (ema_update and learnable_codebook), 'learnable codebook not compatible with EMA update'
+        assert not (exists(vq_bridge) and not learnable_codebook), 'learnable_codebook must be set to True if vq_bridge is passed in'
+        assert not (exists(vq_bridge) and ema_update), 'ema_update must be False if vq_bridge is passed in'
 
         assert 0 <= sync_update_v <= 1.
         assert not (sync_update_v > 0. and not learnable_codebook), 'learnable codebook must be turned on'