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: Implements the vector quantization with straight through approximation of the gradient.
  
  Arguments
  ---------
  z_e_x : torch.Tensor
      Input tensor to be quantized.
  labels : torch.Tensor
      Predicted class for input representations (used for latent space quantization).
  
  Returns
  --------
  Straigth through quantized representation and quantized representation : tuple
  
  Example:
  --------
  >>> inputs = torch.ones(3, 256, 14, 25)
  >>> codebook = VQEmbedding(1024, 256)
  >>> labels = torch.Tensor([1, 0, 2])
  >>> quant, quant_ind = codebook.straight_through(inputs, labels)
  >>> print(quant.shape, quant_ind.shape)
  torch.Size([3, 256, 14, 25]) torch.Size([3, 256, 14, 25])
docstr_tok: ["Implements","the","vector","quantization","with","straight","through","approximation","of","the","gradient",".","Arguments","--","--","--","--","-","z_e_x",":","torch.Tensor","Input","tensor","to","be","quantized",".","labels",":","torch.Tensor","Predicted","class","for","input","representations","(","used","for","latent","space","quantization",")",".","Returns","--","--","--","--","Straigth","through","quantized","representation","and","quantized","representation",":","tuple","Example",":","--","--","--","--",">",">",">","inputs","=","torch.ones","(","3",",","256",",","14",",","25",")",">",">",">","codebook","=","VQEmbedding","(","1024",",","256",")",">",">",">","labels","=","torch.Tensor","(","[","1",",","0",",","2","]",")",">",">",">","quant",",","quant_ind","=","codebook.straight_through","(","inputs",",","labels",")",">",">",">","print","(","quant.shape",",","quant_ind.shape",")","torch.Size","(","[","3",",","256",",","14",",","25","]",")","torch.Size","(","[","3",",","256",",","14",",","25","]",")"]
code: def straight_through(self, z_e_x, labels=None):
      """
      Implements the vector quantization with straight through approximation of the gradient.
  
      Arguments
      ---------
      z_e_x : torch.Tensor
          Input tensor to be quantized.
      labels : torch.Tensor
          Predicted class for input representations (used for latent space quantization).
  
      Returns
      --------
      Straigth through quantized representation and quantized representation : tuple
  
      Example:
      --------
      >>> inputs = torch.ones(3, 256, 14, 25)
      >>> codebook = VQEmbedding(1024, 256)
      >>> labels = torch.Tensor([1, 0, 2])
      >>> quant, quant_ind = codebook.straight_through(inputs, labels)
      >>> print(quant.shape, quant_ind.shape)
      torch.Size([3, 256, 14, 25]) torch.Size([3, 256, 14, 25])
  
      """
      z_e_x_ = z_e_x.permute(0, 2, 3, 1).contiguous()
      z_q_x_, indices = VectorQuantizationStraightThrough.apply(
          z_e_x_,
          self.embedding.weight.detach(),
          labels,
          self.numclasses,
          self.activate_class_partitioning,
          self.shared_keys,
          self.training,
      )
      z_q_x = z_q_x_.permute(0, 3, 1, 2).contiguous()
  
      z_q_x_bar_flatten = torch.index_select(
          self.embedding.weight, dim=0, index=indices
      )
      z_q_x_bar_ = z_q_x_bar_flatten.view_as(z_e_x_)
      z_q_x_bar = z_q_x_bar_.permute(0, 3, 1, 2).contiguous()
  
      return z_q_x, z_q_x_bar
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 3 for calling the VQEmbedding obj's underlying member method code with example usage: obj.straight_through(z_e_x, labels) and returns: z_q_x,  z_q_x_bar