class_docstr: Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly avoids costly matrix
  multiplications and allows for post-hoc remapping of indices.
code: def get_codebook_entry(
      self, indices: torch.LongTensor, shape: Tuple[int, ...]
  ) -> torch.FloatTensor:
      # shape specifying (batch, height, width, channel)
      if self.remap is not None:
          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: torch.FloatTensor = 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: from dataclasses import dataclass
  from typing import Optional, Tuple
  import numpy
  import torch
  import torch.nn
  from ...utils import BaseOutput, is_torch_version
  from ...utils.torch_utils import randn_tensor
  from ..activations import get_activation
  from ..attention_processor import SpatialNorm
  from ..unet_2d_blocks import AutoencoderTinyBlock, UNetMidBlock2D, get_down_block, get_up_block
label_desc: Use image node_id 5 for calling the VectorQuantizer obj's underlying member method code with example usage: obj.get_codebook_entry(indices, shape) and returns: z_q