code: def forward(self, z):
      # reshape z -> (batch, height, width, channel) and flatten
      z = z.permute(0, 2, 3, 1).contiguous()
      z_flattened = z.view(-1, self.emb_dim)
  
      # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
      d = (
          (z_flattened**2).sum(dim=1, keepdim=True)
          + (self.embedding.weight**2).sum(1)
          - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
      )
  
      mean_distance = torch.mean(d)
      # find closest encodings
      # min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
      min_encoding_scores, min_encoding_indices = torch.topk(
          d, 1, dim=1, largest=False
      )
      # [0-1], higher score, higher confidence
      min_encoding_scores = torch.exp(-min_encoding_scores / 10)
  
      min_encodings = torch.zeros(
          min_encoding_indices.shape[0], self.codebook_size
      ).to(z)
      min_encodings.scatter_(1, min_encoding_indices, 1)
  
      # get quantized latent vectors
      z_q = torch.matmul(min_encodings, self.embedding.weight).view(
          z.shape
      )
      # compute loss for embedding
      loss = torch.mean(
          (z_q.detach() - z) ** 2
      ) + self.beta * torch.mean((z_q - z.detach()) ** 2)
      # preserve gradients
      z_q = z + (z_q - z).detach()
  
      # perplexity
      e_mean = torch.mean(min_encodings, dim=0)
      perplexity = torch.exp(
          -torch.sum(e_mean * torch.log(e_mean + 1e-10))
      )
      # reshape back to match original input shape
      z_q = z_q.permute(0, 3, 1, 2).contiguous()
  
      return (
          z_q,
          loss,
          {
              "perplexity": perplexity,
              "min_encodings": min_encodings,
              "min_encoding_indices": min_encoding_indices,
              "min_encoding_scores": min_encoding_scores,
              "mean_distance": mean_distance,
          },
      )
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from basicsr.utils import get_root_logger
  from basicsr.utils.registry import ARCH_REGISTRY
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