class_docstr: credit to @karpathy:
  https://github.com/karpathy/deep-vector-quantization/blob/main/model.py (thanks!)
  Gumbel Softmax trick quantizer
  Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016
  https://arxiv.org/abs/1611.01144
code: def __init__(
      self,
      num_hiddens: int,
      embedding_dim: int,
      n_embed: int,
      straight_through: bool = True,
      kl_weight: float = 5e-4,
      temp_init: float = 1.0,
      remap: Optional[str] = None,
      unknown_index: str = "random",
      loss_key: str = "loss/vq",
  ) -> None:
      super().__init__()
  
      self.loss_key = loss_key
      self.embedding_dim = embedding_dim
      self.n_embed = n_embed
  
      self.straight_through = straight_through
      self.temperature = temp_init
      self.kl_weight = kl_weight
  
      self.proj = nn.Conv2d(num_hiddens, n_embed, 1)
      self.embed = nn.Embedding(n_embed, embedding_dim)
  
      self.remap = remap
      if self.remap is not None:
          self.register_buffer(
              "used", torch.tensor(np.load(self.remap))
          )
          self.re_embed = self.used.shape[0]
      else:
          self.used = None
          self.re_embed = n_embed
      if unknown_index == "extra":
          self.unknown_index = self.re_embed
          self.re_embed = self.re_embed + 1
      else:
          assert unknown_index == "random" or isinstance(
              unknown_index, int
          ), "unknown index needs to be 'random', 'extra' or any integer"
          self.unknown_index = (
              unknown_index  # "random" or "extra" or integer
          )
      if self.remap is not None:
          logpy.info(
              f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
              f"Using {self.unknown_index} for unknown indices."
          )
  
imports: import logging
  from abc import abstractmethod
  from typing import Dict, Iterator, Literal, Optional, Tuple, Union
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from torch import einsum
  from .base import AbstractRegularizer, measure_perplexity
label_desc: Use image node_id 1 to create a new GumbelQuantizer object from inherited base classes: AbstractQuantizer with example: obj = GumbelQuantizer(num_hiddens, embedding_dim, n_embed, straight_through, kl_weight, temp_init, remap, unknown_index, loss_key)