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 forward(
      self,
      z: torch.Tensor,
  ) -> Tuple[torch.Tensor, Dict]:
      do_reshape = z.ndim == 4
      if do_reshape:
          #     # reshape z -> (batch, height, width, channel) and flatten
          z = rearrange(z, "b c h w -> b h w c").contiguous()
  
      else:
          assert (
              z.ndim < 4
          ), "No reshaping strategy for inputs > 4 dimensions defined"
          z = z.contiguous()
  
      z_flattened = z.view(-1, self.e_dim)
      # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
  
      d = (
          torch.sum(z_flattened**2, dim=1, keepdim=True)
          + torch.sum(self.embedding.weight**2, dim=1)
          - 2
          * torch.einsum(
              "bd,dn->bn",
              z_flattened,
              rearrange(self.embedding.weight, "n d -> d n"),
          )
      )
  
      min_encoding_indices = torch.argmin(d, dim=1)
      z_q = self.embedding(min_encoding_indices).view(z.shape)
      loss_dict = {}
      if self.log_perplexity:
          perplexity, cluster_usage = measure_perplexity(
              min_encoding_indices.detach(), self.n_e
          )
          loss_dict.update(
              {"perplexity": perplexity, "cluster_usage": cluster_usage}
          )
  
      # compute loss for embedding
      loss = self.beta * torch.mean(
          (z_q.detach() - z) ** 2
      ) + torch.mean((z_q - z.detach()) ** 2)
      loss_dict[self.loss_key] = loss
  
      # preserve gradients
      z_q = z + (z_q - z).detach()
  
      # reshape back to match original input shape
      if do_reshape:
          z_q = rearrange(z_q, "b h w c -> b c h w").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:
          if do_reshape:
              min_encoding_indices = min_encoding_indices.reshape(
                  z_q.shape[0], z_q.shape[2], z_q.shape[3]
              )
          else:
              min_encoding_indices = rearrange(
                  min_encoding_indices, "(b s) 1 -> b s", b=z_q.shape[0]
              )
  
      loss_dict["min_encoding_indices"] = min_encoding_indices
  
      return z_q, loss_dict
  
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 2 for calling the VectorQuantizer obj's underlying member method code with example usage: obj.forward(z) and returns: z_q,  loss_dict