class_docstr: An implementation for Vector Quantization
class_docstr_tok: ["An","implementation","for","Vector","Quantization"]
docstr: Quantized the input tensor using a fixed codebook and returns
  the corresponding codebook vectors
  
  Parameters
  ----------
  z : Tensor[B x D x T]
  
  Returns
  -------
  Tensor[B x D x T]
      Quantized continuous representation of input
  Tensor[1]
      Commitment loss to train encoder to predict vectors closer to codebook
      entries
  Tensor[1]
      Codebook loss to update the codebook
  Tensor[B x T]
      Codebook indices (quantized discrete representation of input)
  Tensor[B x D x T]
      Projected latents (continuous representation of input before quantization)
docstr_tok: ["Quantized","the","input","tensor","using","a","fixed","codebook","and","returns","the","corresponding","codebook","vectors","Parameters","--","--","--","--","--","z",":","Tensor","[","B","x","D","x","T","]","Returns","--","--","--","-","Tensor","[","B","x","D","x","T","]","Quantized","continuous","representation","of","input","Tensor","[","1","]","Commitment","loss","to","train","encoder","to","predict","vectors","closer","to","codebook","entries","Tensor","[","1","]","Codebook","loss","to","update","the","codebook","Tensor","[","B","x","T","]","Codebook","indices","(","quantized","discrete","representation","of","input",")","Tensor","[","B","x","D","x","T","]","Projected","latents","(","continuous","representation","of","input","before","quantization",")"]
code: def forward(self, z: torch.Tensor):
      """Quantized the input tensor using a fixed codebook and returns
      the corresponding codebook vectors
  
      Parameters
      ----------
      z : Tensor[B x D x T]
  
      Returns
      -------
      Tensor[B x D x T]
          Quantized continuous representation of input
      Tensor[1]
          Commitment loss to train encoder to predict vectors closer to codebook
          entries
      Tensor[1]
          Codebook loss to update the codebook
      Tensor[B x T]
          Codebook indices (quantized discrete representation of input)
      Tensor[B x D x T]
          Projected latents (continuous representation of input before quantization)
      """
  
      # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
      z_e = self.in_proj(z)  # z_e : (B x D x T)
      z_q, indices = self.decode_latents(z_e)
  
      commitment_loss = F.mse_loss(
          z_e, z_q.detach(), reduction="none"
      ).mean([1, 2])
      codebook_loss = F.mse_loss(
          z_q, z_e.detach(), reduction="none"
      ).mean([1, 2])
  
      z_q = (
          z_e + (z_q - z_e).detach()
      )  # noop in forward pass, straight-through gradient estimator in backward pass
  
      z_q = self.out_proj(z_q)
  
      return z_q, commitment_loss, codebook_loss, indices, z_e
  
imports: import math
  from pathlib import Path
  from typing import List, Union
  import logging
  import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.nn.utils.parametrizations import weight_norm
label_desc: Use image node_id 2 for calling the VectorQuantize obj's underlying member method code with example usage: obj.forward(z) and returns: z_q,  commitment_loss,  codebook_loss,  indices,  z_e