class_docstr: This class defines the forward method for vector quantization. As VQ is not differentiable, it returns a RuntimeError in case `.grad()` is called. Refer to `VectorQuantizationStraightThrough` for a straight_through estimation of the gradient for the VQ operation.
docstr: Applies VQ to vectors `input` with `codebook` as VQ dictionary.
  
  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
  --------
  Codebook's indices for quantized representation : torch.Tensor
  
  Example:
  --------
  >>> inputs = torch.ones(3, 14, 25, 256)
  >>> codebook = torch.randn(1024, 256)
  >>> labels = torch.Tensor([1, 0, 2])
  >>> print(VectorQuantization.apply(inputs, codebook, labels).shape)
  torch.Size([3, 14, 25])
docstr_tok: ["Applies","VQ","to","vectors","`","input","`","with","`","codebook","`","as","VQ","dictionary",".","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","--","--","--","--","Codebook","'s","indices","for","quantized","representation",":","torch.Tensor","Example",":","--","--","--","--",">",">",">","inputs","=","torch.ones","(","3",",","14",",","25",",","256",")",">",">",">","codebook","=","torch.randn","(","1024",",","256",")",">",">",">","labels","=","torch.Tensor","(","[","1",",","0",",","2","]",")",">",">",">","print","(","VectorQuantization.apply","(","inputs",",","codebook",",","labels",")",".shape",")","torch.Size","(","[","3",",","14",",","25","]",")"]
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.
  
      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
      --------
      Codebook's indices for quantized representation : torch.Tensor
  
      Example:
      --------
      >>> inputs = torch.ones(3, 14, 25, 256)
      >>> codebook = torch.randn(1024, 256)
      >>> labels = torch.Tensor([1, 0, 2])
      >>> print(VectorQuantization.apply(inputs, codebook, labels).shape)
      torch.Size([3, 14, 25])
      """
      with torch.no_grad():
          embedding_size = codebook.size(1)
          inputs_size = inputs.size()
          inputs_flatten = inputs.view(-1, embedding_size)
  
          labels_expanded = labels.reshape(-1, 1, 1).repeat(
              1, inputs_size[1], inputs_size[2]
          )
          labels_flatten = labels_expanded.reshape(-1)
          irrelevant_regions = get_irrelevant_regions(
              labels_flatten,
              codebook.shape[0],
              num_classes,
              N_shared=shared_keys,
              stage="TRAIN" if training else "VALID",
          )
  
          codebook_sqr = torch.sum(codebook**2, dim=1)
          inputs_sqr = torch.sum(
              inputs_flatten**2, dim=1, keepdim=True
          )
  
          # Compute the distances to the codebook
          distances = torch.addmm(
              codebook_sqr + inputs_sqr,
              inputs_flatten,
              codebook.t(),
              alpha=-2.0,
              beta=1.0,
          )
  
          # intervene and boost the distances for irrelevant codes
          if activate_class_partitioning:
              distances[irrelevant_regions] = torch.inf
  
          _, indices_flatten = torch.min(distances, dim=1)
          indices = indices_flatten.view(*inputs_size[:-1])
          ctx.mark_non_differentiable(indices)
  
          return indices
  
imports: import torch
  import torch.nn
  from torch.autograd import Function
label_desc: Use image node_id 1 for calling the VectorQuantization 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: indices