class_docstr: Implements VQ Dictionary. Wraps `VectorQuantization` and `VectorQuantizationStraightThrough`. For more details refer to the specific class.
  
  Arguments
  ---------
  K : int
      Number of elements of VQ dictionary.
  D : int
      Dimensionality of VQ vectors.
  num_classes : int
      Number of possible classes
  activate_class_partitioning : bool
      `True` if latent space should be quantized for different classes.
  shared_keys : int
      Number of shared keys among classes.
docstr: Wraps VectorQuantization. Computes VQ-dictionary indices for input quantization. Note that this forward step is not differentiable.
  
  Arguments
  ---------
  z_e_x : torch.Tensor
      Input tensor to be quantized.
  
  Returns
  --------
  Codebook's indices for quantized representation : torch.Tensor
  
  Example:
  --------
  >>> inputs = torch.ones(3, 256, 14, 25)
  >>> codebook = VQEmbedding(1024, 256)
  >>> labels = torch.Tensor([1, 0, 2])
  >>> print(codebook(inputs, labels).shape)
  torch.Size([3, 14, 25])
docstr_tok: ["Wraps","VectorQuantization",".","Computes","VQ-dictionary","indices","for","input","quantization",".","Note","that","this","forward","step","is","not","differentiable",".","Arguments","--","--","--","--","-","z_e_x",":","torch.Tensor","Input","tensor","to","be","quantized",".","Returns","--","--","--","--","Codebook","'s","indices","for","quantized","representation",":","torch.Tensor","Example",":","--","--","--","--",">",">",">","inputs","=","torch.ones","(","3",",","256",",","14",",","25",")",">",">",">","codebook","=","VQEmbedding","(","1024",",","256",")",">",">",">","labels","=","torch.Tensor","(","[","1",",","0",",","2","]",")",">",">",">","print","(","codebook","(","inputs",",","labels",")",".shape",")","torch.Size","(","[","3",",","14",",","25","]",")"]
code: def forward(self, z_e_x, labels=None):
      """
      Wraps VectorQuantization. Computes VQ-dictionary indices for input quantization. Note that this forward step is not differentiable.
  
      Arguments
      ---------
      z_e_x : torch.Tensor
          Input tensor to be quantized.
  
      Returns
      --------
      Codebook's indices for quantized representation : torch.Tensor
  
      Example:
      --------
      >>> inputs = torch.ones(3, 256, 14, 25)
      >>> codebook = VQEmbedding(1024, 256)
      >>> labels = torch.Tensor([1, 0, 2])
      >>> print(codebook(inputs, labels).shape)
      torch.Size([3, 14, 25])
      """
      z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
      latents = VectorQuantization.apply(
          z_e_x_, self.embedding.weight, labels
      )
      return latents
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 2 for calling the VQEmbedding obj's underlying member method code with example usage: obj.forward(z_e_x, labels) and returns: latents