This proposition is both a refacto of a few parts of the LatentQuantizer, and the resolution of some bugs. The changes have been checked consistently by the addition of a dedicated test file which covers (almost) all the code in LatentQuantize.
The main goal was to improve the code by simplifying as much as possible the multiple branches: same levels or not, images or video or not, etc. The code now deals the kind of input the same way, as it does for the way levels are provided (by a list or int).
I will happily answer any question about my PR :)
Here is a list of the proposed modifications and their details:
Addition
Add a test file to check the refacto.
Add a test dedicated to "same level" case of LatentQuantization --> uncover an error in the quantize method.
Add test for 'int' case.
Add: more tests to check time series, images and video features.
Add a 'codebook_dim' parameter to the init function, default to -1. Used only if 'levels' is an int. The positivity of the value is checked to ensure the construction of a list of codebook_dim times the value 'levels'.
Add a test to check the previous point.
Fixes
Fix: The dim attribute is not optional: change of the docstring, remove the use of 'default' computation.
Fix: same level case, by simply elimintating the '_equal_levels' variable, and ensuring that both cases are treated the same way. In particular, 'values_per_latent' parameter has the same type and shape (that is, a ParameterList)
Fix: For time series, the 'features' dimension should match the 'dim' parameter of the LatentQuantizer.
Fix: the variable 'in_place_codebook_optimizer' passed in init was immediately passed to 'None' as an attribute.
Format
Format and lint of code using Ruff
Removes
Remove: helper function 'exists' and 'defaults'. They added a level of complexity without adding lisibility (in particular 'default' and its absence of docstring)
This proposition is both a refacto of a few parts of the LatentQuantizer, and the resolution of some bugs. The changes have been checked consistently by the addition of a dedicated test file which covers (almost) all the code in LatentQuantize.
The main goal was to improve the code by simplifying as much as possible the multiple branches: same levels or not, images or video or not, etc. The code now deals the kind of input the same way, as it does for the way levels are provided (by a list or int).
I will happily answer any question about my PR :)
Here is a list of the proposed modifications and their details:
Addition
Fixes
Format
Removes