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 get_codebook_entry(self, indices, shape):
      # TODO: shape not yet optional
      b, h, w, c = shape
      assert b * h * w == indices.shape[0]
      indices = rearrange(indices, "(b h w) -> b h w", b=b, h=h, w=w)
      if self.remap is not None:
          indices = self.unmap_to_all(indices)
      one_hot = (
          F.one_hot(indices, num_classes=self.n_embed)
          .permute(0, 3, 1, 2)
          .float()
      )
      z_q = einsum(
          "b n h w, n d -> b d h w", one_hot, self.embed.weight
      )
      return z_q
  
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 3 for calling the GumbelQuantizer obj's underlying member method code with example usage: obj.get_codebook_entry(indices, shape) and returns: z_q