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 __init__(
      self,
      n_e: int,
      e_dim: int,
      beta: float = 0.25,
      remap: Optional[str] = None,
      unknown_index: str = "random",
      sane_index_shape: bool = False,
      log_perplexity: bool = False,
      embedding_weight_norm: bool = False,
      loss_key: str = "loss/vq",
  ):
      super().__init__()
      self.n_e = n_e
      self.e_dim = e_dim
      self.beta = beta
      self.loss_key = loss_key
  
      if not embedding_weight_norm:
          self.embedding = nn.Embedding(self.n_e, self.e_dim)
          self.embedding.weight.data.uniform_(
              -1.0 / self.n_e, 1.0 / self.n_e
          )
      else:
          self.embedding = torch.nn.utils.weight_norm(
              nn.Embedding(self.n_e, self.e_dim), dim=1
          )
  
      self.remap = remap
      if self.remap is not None:
          self.register_buffer(
              "used", torch.tensor(np.load(self.remap))
          )
          self.re_embed = self.used.shape[0]
      else:
          self.used = None
          self.re_embed = n_e
      if unknown_index == "extra":
          self.unknown_index = self.re_embed
          self.re_embed = self.re_embed + 1
      else:
          assert unknown_index == "random" or isinstance(
              unknown_index, int
          ), "unknown index needs to be 'random', 'extra' or any integer"
          self.unknown_index = (
              unknown_index  # "random" or "extra" or integer
          )
      if self.remap is not None:
          logpy.info(
              f"Remapping {self.n_e} indices to {self.re_embed} indices. "
              f"Using {self.unknown_index} for unknown indices."
          )
  
      self.sane_index_shape = sane_index_shape
      self.log_perplexity = log_perplexity
  
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 1 to create a new VectorQuantizer object from inherited base classes: AbstractQuantizer with example: obj = VectorQuantizer(n_e, e_dim, beta, remap, unknown_index, sane_index_shape, log_perplexity, embedding_weight_norm, loss_key)