docstr: Vector quantization using gumbel softmax
  
  Args:
      dim: input dimension (channels)
      num_vars: number of quantized vectors per group
      temp: temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor)
      groups: number of groups for vector quantization
      combine_groups: whether to use the vectors for all groups
      vq_dim: dimensionality of the resulting quantized vector
      time_first: if true, expect input in BxTxC format, otherwise in BxCxT
      activation: what activation to use (should be a module). this is only used if weight_proj_depth is > 1
      weight_proj_depth: number of layers (with activation in between) to project input before computing logits
      weight_proj_factor: this is used only if weight_proj_depth is > 1. scales the inner dimensionality of
                          projections by this factor
docstr_tok: ["Vector","quantization","using","gumbel","softmax","Args",":","dim",":","input","dimension","(","channels",")","num_vars",":","number","of","quantized","vectors","per","group","temp",":","temperature","for","training",".","this","should","be","a","tuple","of","3","elements",":","(","start",",","stop",",","decay","factor",")","groups",":","number","of","groups","for","vector","quantization","combine_groups",":","whether","to","use","the","vectors","for","all","groups","vq_dim",":","dimensionality","of","the","resulting","quantized","vector","time_first",":","if","true",",","expect","input","in","BxTxC","format",",","otherwise","in","BxCxT","activation",":","what","activation","to","use","(","should","be","a","module",")",".","this","is","only","used","if","weight_proj_depth","is",">","1","weight_proj_depth",":","number","of","layers","(","with","activation","in","between",")","to","project","input","before","computing","logits","weight_proj_factor",":","this","is","used","only","if","weight_proj_depth","is",">","1.","scales","the","inner","dimensionality","of","projections","by","this","factor"]
code: def __init__(
      self,
      dim,
      num_vars,
      temp,
      groups,
      combine_groups,
      vq_dim,
      time_first,
      activation="gelu",
      weight_proj_depth=1,
      weight_proj_factor=1,
  ):
      """Vector quantization using gumbel softmax
  
      Args:
          dim: input dimension (channels)
          num_vars: number of quantized vectors per group
          temp: temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor)
          groups: number of groups for vector quantization
          combine_groups: whether to use the vectors for all groups
          vq_dim: dimensionality of the resulting quantized vector
          time_first: if true, expect input in BxTxC format, otherwise in BxCxT
          activation: what activation to use (should be a module). this is only used if weight_proj_depth is > 1
          weight_proj_depth: number of layers (with activation in between) to project input before computing logits
          weight_proj_factor: this is used only if weight_proj_depth is > 1. scales the inner dimensionality of
                              projections by this factor
      """
      super().__init__()
  
      self.groups = groups
      self.combine_groups = combine_groups
      self.input_dim = dim
      self.num_vars = num_vars
      self.time_first = time_first
  
      assert (
          vq_dim % groups == 0
      ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation"
  
      var_dim = vq_dim // groups
      num_groups = groups if not combine_groups else 1
  
      self.vars = nn.Parameter(
          torch.FloatTensor(1, num_groups * num_vars, var_dim)
      )
      nn.init.uniform_(self.vars)
  
      if weight_proj_depth > 1:
          activation = jasper_activations["gelu"]
  
          def block(input_dim, output_dim):
              return nn.Sequential(
                  nn.Linear(input_dim, output_dim), activation
              )
  
          inner_dim = self.input_dim * weight_proj_factor
          self.weight_proj = nn.Sequential(
              *[
                  block(
                      self.input_dim if i == 0 else inner_dim, inner_dim
                  )
                  for i in range(weight_proj_depth - 1)
              ],
              nn.Linear(inner_dim, groups * num_vars),
          )
      else:
          self.weight_proj = nn.Linear(
              self.input_dim, groups * num_vars
          )
          nn.init.normal_(self.weight_proj.weight, mean=0, std=1)
          nn.init.zeros_(self.weight_proj.bias)
  
      assert (
          len(temp) == 3
      ), "Quantize temperature should be a tuple of 3 elements: (start, stop, decay factor)"
  
      self.max_temp, self.min_temp, self.temp_decay = temp
      self.curr_temp = self.max_temp
      self.codebook_indices = None
  
file_path: NeMo/nemo/collections/asr/parts/submodules/ssl_quantizers.py
imports: import torch
  import torch.nn.functional
  from torch import nn
  from nemo.collections.asr.parts.submodules.jasper import jasper_activations
  from nemo.core import NeuralModule
  from nemo.core.neural_types import EncodedRepresentation, LossType, NeuralType
label_desc: Use image node_id 1 to create a new GumbelVectorQuantizer object from inherited base classes: NeuralModule with example: obj = GumbelVectorQuantizer(dim, num_vars, temp, groups, combine_groups, vq_dim, time_first, activation, weight_proj_depth, weight_proj_factor)