class_docstr: Vector quantization using gumbel softmax. Copied from fairseq implementation.
  Arguments
  ---------
      input_dim: int
          Input dimension (channels).
      num_vars: int
          Number of quantized vectors per group.
      temp_tuple: float
          Temperature for training. this should be a tuple of 3 elements: (start, stop, decay factor).
      groups: int
          Number of groups for vector quantization.
      vq_dim: int
          Dimensionality of the resulting quantized vector.
  
  Example
  -------
  >>> quantiser = GumbelVectorQuantizer(128, 100, (2.0, 0.25, 0.999995,), 2, 50 )
  >>> inputs = torch.rand(10, 12, 128)
  >>> output = quantiser(inputs)
  >>> output["x"].shape
  torch.Size([10, 12, 50])
code: def __init__(self, input_dim, num_vars, temp_tuple, groups, vq_dim):
      super().__init__()
  
      self.groups = groups
      self.input_dim = input_dim
      self.num_vars = num_vars
      self.vq_dim = vq_dim
  
      assert (
          vq_dim % groups == 0
      ), f"dim {vq_dim} must be divisible by groups {groups} for concatenation"
  
      var_dim = vq_dim // groups
  
      self.vars = nn.Parameter(
          torch.FloatTensor(1, groups * num_vars, var_dim)
      )
      nn.init.uniform_(self.vars)
  
      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_tuple) == 3, temp_tuple
  
      self.max_temp, self.min_temp, self.temp_decay = temp_tuple
      self.curr_temp = self.max_temp
      self.max_ent = nn.Parameter(
          torch.log(torch.tensor(float(self.num_vars * self.groups))),
          requires_grad=False,
      )
  
imports: import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 1 to create a new GumbelVectorQuantizer object from inherited base classes: nn with example: obj = GumbelVectorQuantizer(input_dim, num_vars, temp_tuple, groups, vq_dim)