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 forward(
      self,
      z: torch.Tensor,
      temp: Optional[float] = None,
      return_logits: bool = False,
  ) -> Tuple[torch.Tensor, Dict]:
      # force hard = True when we are in eval mode, as we must quantize.
      # actually, always true seems to work
      hard = self.straight_through if self.training else True
      temp = self.temperature if temp is None else temp
      out_dict = {}
      logits = self.proj(z)
      if self.remap is not None:
          # continue only with used logits
          full_zeros = torch.zeros_like(logits)
          logits = logits[:, self.used, ...]
  
      soft_one_hot = F.gumbel_softmax(
          logits, tau=temp, dim=1, hard=hard
      )
      if self.remap is not None:
          # go back to all entries but unused set to zero
          full_zeros[:, self.used, ...] = soft_one_hot
          soft_one_hot = full_zeros
      z_q = einsum(
          "b n h w, n d -> b d h w", soft_one_hot, self.embed.weight
      )
  
      # + kl divergence to the prior loss
      qy = F.softmax(logits, dim=1)
      diff = (
          self.kl_weight
          * torch.sum(
              qy * torch.log(qy * self.n_embed + 1e-10), dim=1
          ).mean()
      )
      out_dict[self.loss_key] = diff
  
      ind = soft_one_hot.argmax(dim=1)
      out_dict["indices"] = ind
      if self.remap is not None:
          ind = self.remap_to_used(ind)
  
      if return_logits:
          out_dict["logits"] = logits
  
      return z_q, out_dict
  
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 2 for calling the GumbelQuantizer obj's underlying member method code with example usage: obj.forward(z, temp, return_logits) and returns: z_q,  out_dict