class_docstr: An implementation for Vector Quantization
class_docstr_tok: ["An","implementation","for","Vector","Quantization"]
docstr: Implementation of VQ similar to Karpathy's repo:
  https://github.com/karpathy/deep-vector-quantization
  Additionally uses following tricks from Improved VQGAN
  (https://arxiv.org/pdf/2110.04627.pdf):
      1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
          for improved codebook usage
      2. l2-normalized codes: Converts euclidean distance to cosine similarity which
          improves training stability
docstr_tok: ["Implementation","of","VQ","similar","to","Karpathy","'s","repo",":","https",":","\/\/github.com\/karpathy\/deep-vector-quantization","Additionally","uses","following","tricks","from","Improved","VQGAN","(","https",":","\/\/arxiv.org\/pdf\/2110.04627.pdf",")",":","1",".","Factorized","codes",":","Perform","nearest","neighbor","lookup","in","low-dimensional","space","for","improved","codebook","usage","2.","l2-normalized","codes",":","Converts","euclidean","distance","to","cosine","similarity","which","improves","training","stability"]
code: def __init__(
      self, input_dim: int, codebook_size: int, codebook_dim: int
  ):
      """
      Implementation of VQ similar to Karpathy's repo:
      https://github.com/karpathy/deep-vector-quantization
      Additionally uses following tricks from Improved VQGAN
      (https://arxiv.org/pdf/2110.04627.pdf):
          1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
              for improved codebook usage
          2. l2-normalized codes: Converts euclidean distance to cosine similarity which
              improves training stability
      """
      super().__init__()
      self.codebook_size = codebook_size
      self.codebook_dim = codebook_dim
  
      self.in_proj = WNConv1d(input_dim, codebook_dim, kernel_size=1)
      self.out_proj = WNConv1d(codebook_dim, input_dim, kernel_size=1)
      self.codebook = nn.Embedding(codebook_size, codebook_dim)
  
imports: import math
  from pathlib import Path
  from typing import List, Union
  import logging
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.utils.parametrizations import weight_norm
label_desc: Use image node_id 1 to create a new VectorQuantize object from inherited base classes: nn with example: obj = VectorQuantize(input_dim, codebook_size, codebook_dim)