assert that residual quantizers never have projections except before the residual quantizing

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.11.5',
+  version = '1.11.6',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/finite_scalar_quantization.py

@@ -68,8 +68,11 @@ def __init__(
         self.keep_num_codebooks_dim = keep_num_codebooks_dim
 
         self.dim = default(dim, len(_levels) * num_codebooks)
-        self.project_in = nn.Linear(self.dim, effective_codebook_dim) if self.dim != effective_codebook_dim else nn.Identity()
-        self.project_out = nn.Linear(effective_codebook_dim, self.dim) if self.dim != effective_codebook_dim else nn.Identity()
+
+        has_projections = self.dim != effective_codebook_dim
+        self.project_in = nn.Linear(self.dim, effective_codebook_dim) if has_projections else nn.Identity()
+        self.project_out = nn.Linear(effective_codebook_dim, self.dim) if has_projections else nn.Identity()
+        self.has_projections = has_projections
 
         self.codebook_size = self._levels.prod().item()
 

vector_quantize_pytorch/lookup_free_quantization.py

@@ -84,8 +84,10 @@ def __init__(
         codebook_dims = codebook_dim * num_codebooks
         dim = default(dim, codebook_dims)
 
-        self.project_in = nn.Linear(dim, codebook_dims) if dim != codebook_dims else nn.Identity()
-        self.project_out = nn.Linear(codebook_dims, dim) if dim != codebook_dims else nn.Identity()
+        has_projections = dim != codebook_dims
+        self.project_in = nn.Linear(dim, codebook_dims) if has_projections else nn.Identity()
+        self.project_out = nn.Linear(codebook_dims, dim) if has_projections else nn.Identity()
+        self.has_projections = has_projections
 
         self.dim = dim
         self.codebook_dim = codebook_dim

vector_quantize_pytorch/residual_fsq.py

@@ -47,6 +47,7 @@ def __init__(
         requires_projection = codebook_dim != dim
         self.project_in = nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity()
         self.project_out = nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
+        self.has_projections = requires_projection
 
         self.num_quantizers = num_quantizers
 
@@ -68,6 +69,8 @@ def __init__(
 
             self.layers.append(fsq)
 
+        assert all([not fsq.has_projections for fsq in self.layers])
+
         self.codebook_size = self.layers[0].codebook_size
 
         self.register_buffer('scales', torch.stack(scales), persistent = False)

vector_quantize_pytorch/residual_lfq.py

@@ -45,6 +45,7 @@ def __init__(
         requires_projection = codebook_dim != dim
         self.project_in = nn.Linear(dim, codebook_dim) if requires_projection else nn.Identity()
         self.project_out = nn.Linear(codebook_dim, dim) if requires_projection else nn.Identity()
+        self.has_projections = requires_projection
 
         self.num_quantizers = num_quantizers
 
@@ -61,6 +62,8 @@ def __init__(
 
             self.layers.append(lfq)
 
+        assert all([not lfq.has_projections for lfq in self.layers])
+
         self.quantize_dropout = quantize_dropout and num_quantizers > 1
 
         assert quantize_dropout_cutoff_index >= 0

vector_quantize_pytorch/residual_vq.py

@@ -47,12 +47,15 @@ def __init__(
         requires_projection = codebook_input_dim != dim
         self.project_in = nn.Linear(dim, codebook_input_dim) if requires_projection else nn.Identity()
         self.project_out = nn.Linear(codebook_input_dim, dim) if requires_projection else nn.Identity()
+        self.has_projections = requires_projection
 
         self.num_quantizers = num_quantizers
 
         self.accept_image_fmap = accept_image_fmap
         self.layers = nn.ModuleList([VectorQuantize(dim = codebook_dim, codebook_dim = codebook_dim, accept_image_fmap = accept_image_fmap, **kwargs) for _ in range(num_quantizers)])
 
+        assert all([not vq.has_projections for vq in self.layers])
+
         self.quantize_dropout = quantize_dropout and num_quantizers > 1
 
         assert quantize_dropout_cutoff_index >= 0

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -732,6 +732,8 @@ def __init__(
         self.project_in = nn.Linear(dim, codebook_input_dim) if requires_projection else nn.Identity()
         self.project_out = nn.Linear(codebook_input_dim, dim) if requires_projection else nn.Identity()
 
+        self.has_projections = requires_projection
+
         self.eps = eps
         self.commitment_weight = commitment_weight
         self.commitment_use_cross_entropy_loss = commitment_use_cross_entropy_loss # whether to use cross entropy loss to codebook as commitment loss