"""VQ-VAE module."""