Here, we will implement a custom layer to encapsulate the vector quantizer logic, which is the central component of VQ-VAEs. Consider an output from the encoder, with shape (batch_size, height, width, num_channels). The vector quantizer will first flatten this output, only keeping the num_channels dimension intact. So, the shape would become (batch_size * height * width, num_channels). The rationale behind this is to treat the total number of channels as the space for the latent embeddings.