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 forward(
      self, z: torch.FloatTensor
  ) -> Tuple[torch.FloatTensor, torch.FloatTensor, Tuple]:
      # reshape z -> (batch, height, width, channel) and flatten
      z = z.permute(0, 2, 3, 1).contiguous()
      z_flattened = z.view(-1, self.vq_embed_dim)
  
      # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
      min_encoding_indices = torch.argmin(
          torch.cdist(z_flattened, self.embedding.weight), dim=1
      )
  
      z_q = self.embedding(min_encoding_indices).view(z.shape)
      perplexity = None
      min_encodings = None
  
      # compute loss for embedding
      if not self.legacy:
          loss = self.beta * torch.mean(
              (z_q.detach() - z) ** 2
          ) + torch.mean((z_q - z.detach()) ** 2)
      else:
          loss = torch.mean(
              (z_q.detach() - z) ** 2
          ) + self.beta * torch.mean((z_q - z.detach()) ** 2)
  
      # preserve gradients
      z_q: torch.FloatTensor = z + (z_q - z).detach()
  
      # reshape back to match original input shape
      z_q = z_q.permute(0, 3, 1, 2).contiguous()
  
      if self.remap is not None:
          min_encoding_indices = min_encoding_indices.reshape(
              z.shape[0], -1
          )  # add batch axis
          min_encoding_indices = self.remap_to_used(
              min_encoding_indices
          )
          min_encoding_indices = min_encoding_indices.reshape(
              -1, 1
          )  # flatten
  
      if self.sane_index_shape:
          min_encoding_indices = min_encoding_indices.reshape(
              z_q.shape[0], z_q.shape[2], z_q.shape[3]
          )
  
      return (
          z_q,
          loss,
          (perplexity, min_encodings, min_encoding_indices),
      )
  
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 4 for calling the VectorQuantizer obj's underlying member method code with example usage: obj.forward(z) and returns: z_q,  loss,  perplexity,  min_encodings,  min_encoding_indices