code: def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, Dict]:
      # reshape z -> (batch, height, width, channel) and flatten
      # z, 'b c h w -> b h w c'
      z = rearrange(z, "b c h w -> b h w c")
      z_flattened = z.reshape(-1, self.codebook_dim)
  
      # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
      d = (
          z_flattened.pow(2).sum(dim=1, keepdim=True)
          + self.embedding.weight.pow(2).sum(dim=1)
          - 2
          * torch.einsum(
              "bd,nd->bn", z_flattened, self.embedding.weight
          )
      )  # 'n d -> d n'
  
      encoding_indices = torch.argmin(d, dim=1)
  
      z_q = self.embedding(encoding_indices).view(z.shape)
      encodings = F.one_hot(encoding_indices, self.num_tokens).type(
          z.dtype
      )
      avg_probs = torch.mean(encodings, dim=0)
      perplexity = torch.exp(
          -torch.sum(avg_probs * torch.log(avg_probs + 1e-10))
      )
  
      if self.training and self.embedding.update:
          # EMA cluster size
          encodings_sum = encodings.sum(0)
          self.embedding.cluster_size_ema_update(encodings_sum)
          # EMA embedding average
          embed_sum = encodings.transpose(0, 1) @ z_flattened
          self.embedding.embed_avg_ema_update(embed_sum)
          # normalize embed_avg and update weight
          self.embedding.weight_update(self.num_tokens)
  
      # compute loss for embedding
      loss = self.beta * F.mse_loss(z_q.detach(), z)
  
      # preserve gradients
      z_q = z + (z_q - z).detach()
  
      # reshape back to match original input shape
      # z_q, 'b h w c -> b c h w'
      z_q = rearrange(z_q, "b h w c -> b c h w")
  
      out_dict = {
          self.loss_key: loss,
          "encodings": encodings,
          "encoding_indices": encoding_indices,
          "perplexity": perplexity,
      }
  
      return z_q, out_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 EMAVectorQuantizer obj's underlying member method code with example usage: obj.forward(z) and returns: z_q,  out_dict