class_docstr: ____________________________________________
  Discretization bottleneck part of the VQ-VAE.
  Inputs:
  - n_e : number of embeddings
  - e_dim : dimension of embedding
  - beta : commitment cost used in loss term,
      beta * ||z_e(x)-sg[e]||^2
  _____________________________________________
code: def get_codebook_entry(
      self,
      indices: torch.Tensor,
      shape: Optional[Tuple[int, ...]] = None,
  ) -> torch.Tensor:
      # shape specifying (batch, height, width, channel)
      if self.remap is not None:
          assert shape is not None, "Need to give shape for remap"
          indices = indices.reshape(shape[0], -1)  # add batch axis
          indices = self.unmap_to_all(indices)
          indices = indices.reshape(-1)  # flatten again
  
      # get quantized latent vectors
      z_q = self.embedding(indices)
  
      if shape is not None:
          z_q = z_q.view(shape)
          # reshape back to match original input shape
          z_q = z_q.permute(0, 3, 1, 2).contiguous()
  
      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 VectorQuantizer obj's underlying member method code with example usage: obj.get_codebook_entry(indices, shape) and returns: z_q