class_docstr: This class defines the forward method for vector quantization. As VQ is not differentiable, it approximates the gradient of the VQ as in https://arxiv.org/abs/1711.00937.
docstr: Applies VQ to vectors `input` with `codebook` as VQ dictionary and estimates gradients with a
  Straight-Through (id) approximation of the quantization steps.
  
  Arguments
  ---------
  inputs : torch.Tensor
      Hidden representations to quantize. Expected shape is `torch.Size([B, W, H, C])`.
  codebook : torch.Tensor
      VQ-dictionary for quantization. Expected shape of `torch.Size([K, C])` with K dictionary elements.
  labels : torch.Tensor
      Classification labels. Used to define irrelevant regions and divide the latent space based on predicted class. Shape should be `torch.Size([B])`.
  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.
  training : bool
      `True` if stage is TRAIN.
  
  Returns
  --------
  Quantized representation and codebook's indices for quantized representation : tuple
  
  Example:
  --------
  >>> inputs = torch.ones(3, 14, 25, 256)
  >>> codebook = torch.randn(1024, 256)
  >>> labels = torch.Tensor([1, 0, 2])
  >>> quant, quant_ind = VectorQuantizationStraightThrough.apply(inputs, codebook, labels)
  >>> print(quant.shape, quant_ind.shape)
  torch.Size([3, 14, 25, 256]) torch.Size([1050])
docstr_tok: ["Applies","VQ","to","vectors","`","input","`","with","`","codebook","`","as","VQ","dictionary","and","estimates","gradients","with","a","Straight-Through","(","id",")","approximation","of","the","quantization","steps",".","Arguments","--","--","--","--","-","inputs",":","torch.Tensor","Hidden","representations","to","quantize",".","Expected","shape","is","`","torch.Size","(","[","B",",","W",",","H",",","C","]",")","`",".","codebook",":","torch.Tensor","VQ-dictionary","for","quantization",".","Expected","shape","of","`","torch.Size","(","[","K",",","C","]",")","`","with","K","dictionary","elements",".","labels",":","torch.Tensor","Classification","labels",".","Used","to","define","irrelevant","regions","and","divide","the","latent","space","based","on","predicted","class",".","Shape","should","be","`","torch.Size","(","[","B","]",")","`",".","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",".","training",":","bool","`","True","`","if","stage","is","TRAIN",".","Returns","--","--","--","--","Quantized","representation","and","codebook","'s","indices","for","quantized","representation",":","tuple","Example",":","--","--","--","--",">",">",">","inputs","=","torch.ones","(","3",",","14",",","25",",","256",")",">",">",">","codebook","=","torch.randn","(","1024",",","256",")",">",">",">","labels","=","torch.Tensor","(","[","1",",","0",",","2","]",")",">",">",">","quant",",","quant_ind","=","VectorQuantizationStraightThrough.apply","(","inputs",",","codebook",",","labels",")",">",">",">","print","(","quant.shape",",","quant_ind.shape",")","torch.Size","(","[","3",",","14",",","25",",","256","]",")","torch.Size","(","[","1050","]",")"]
code: def forward(
      ctx,
      inputs,
      codebook,
      labels=None,
      num_classes=10,
      activate_class_partitioning=True,
      shared_keys=10,
      training=True,
  ):
      """
      Applies VQ to vectors `input` with `codebook` as VQ dictionary and estimates gradients with a
      Straight-Through (id) approximation of the quantization steps.
  
      Arguments
      ---------
      inputs : torch.Tensor
          Hidden representations to quantize. Expected shape is `torch.Size([B, W, H, C])`.
      codebook : torch.Tensor
          VQ-dictionary for quantization. Expected shape of `torch.Size([K, C])` with K dictionary elements.
      labels : torch.Tensor
          Classification labels. Used to define irrelevant regions and divide the latent space based on predicted class. Shape should be `torch.Size([B])`.
      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.
      training : bool
          `True` if stage is TRAIN.
  
      Returns
      --------
      Quantized representation and codebook's indices for quantized representation : tuple
  
      Example:
      --------
      >>> inputs = torch.ones(3, 14, 25, 256)
      >>> codebook = torch.randn(1024, 256)
      >>> labels = torch.Tensor([1, 0, 2])
      >>> quant, quant_ind = VectorQuantizationStraightThrough.apply(inputs, codebook, labels)
      >>> print(quant.shape, quant_ind.shape)
      torch.Size([3, 14, 25, 256]) torch.Size([1050])
      """
      indices = VectorQuantization.apply(
          inputs,
          codebook,
          labels,
          num_classes,
          activate_class_partitioning,
          shared_keys,
          training,
      )
      indices_flatten = indices.view(-1)
      ctx.save_for_backward(indices_flatten, codebook)
      ctx.mark_non_differentiable(indices_flatten)
  
      codes_flatten = torch.index_select(
          codebook, dim=0, index=indices_flatten
      )
      codes = codes_flatten.view_as(inputs)
  
      return (codes, indices_flatten)
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 1 for calling the VectorQuantizationStraightThrough obj's underlying member method code with example usage: obj.forward(ctx, inputs, codebook, labels, num_classes, activate_class_partitioning, shared_keys, training) and returns: codes,  indices_flatten