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.
class_docstr_tok: ["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):
      # reshape z -> (batch, height, width, channel) and flatten
      z = z.transpose([0, 2, 3, 1])
      z_flattened = z.reshape([-1, self.e_dim])
      # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
  
      d = (
          paddle.sum(z_flattened**2, axis=1, keepdim=True)
          + paddle.sum(self.embedding.weight**2, axis=1)
          - 2
          * paddle.einsum(
              "bd,dn->bn", z_flattened, self.embedding.weight.t()
          )
      )
  
      min_encoding_indices = paddle.argmin(d, axis=1)
      z_q = self.embedding(min_encoding_indices).reshape(z.shape)
      perplexity = None
      min_encodings = None
  
      # compute loss for embedding
      if not self.legacy:
          loss = self.beta * paddle.mean(
              (z_q.detach() - z) ** 2
          ) + paddle.mean((z_q - z.detach()) ** 2)
      else:
          loss = paddle.mean(
              (z_q.detach() - z) ** 2
          ) + self.beta * paddle.mean((z_q - z.detach()) ** 2)
  
      # preserve gradients
      z_q = z + (z_q - z).detach()
  
      # reshape back to match original input shape
      z_q = z_q.transpose([0, 3, 1, 2])
  
      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: import numpy
  import paddle
  import paddle.nn
  from ..configuration_utils import ConfigMixin
  from ..configuration_utils import register_to_config
  from .unet_blocks import get_down_block
  from .unet_blocks import get_up_block
  from .unet_blocks import UNetMidBlock2D
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