class_docstr: Introduced in SoundStream: An end2end neural audio codec
  https://arxiv.org/abs/2107.03312
  
  
  Example
  -------
  Using a pretrained RVQ unit.
  
  >>> dac = DAC(load_pretrained=True, model_type="44KHz", model_bitrate="8kbps", tag="latest")
  >>> quantizer = dac.quantizer
  >>> continuous_embeddings = torch.randn(1, 1024, 100) # Example shape: [Batch, Channels, Time]
  >>> discrete_embeddings, codes, _, _, _ = quantizer(continuous_embeddings)
class_docstr_tok: ["Introduced","in","SoundStream",":","An","end2end","neural","audio","codec","https",":","\/\/arxiv.org\/abs\/2107.03312","Example","--","--","--","-","Using","a","pretrained","RVQ","unit",".",">",">",">","dac","=","DAC","(","load_pretrained=True",",","model_type=","''","44KHz","''",",","model_bitrate=","''","8kbps","''",",","tag=","''","latest","''",")",">",">",">","quantizer","=","dac.quantizer",">",">",">","continuous_embeddings","=","torch.randn","(","1",",","1024",",","100",")","#","Example","shape",":","[","Batch",",","Channels",",","Time","]",">",">",">","discrete_embeddings",",","codes",",","_",",","_",",","_","=","quantizer","(","continuous_embeddings",")"]
docstr: Quantized the input tensor using a fixed set of `n` codebooks and returns
  the corresponding codebook vectors
  Parameters
  ----------
  z : Tensor[B x D x T]
  n_quantizers : int, optional
      No. of quantizers to use
      (n_quantizers < self.n_codebooks ex: for quantizer dropout)
      Note: if `self.quantizer_dropout` is True, this argument is ignored
          when in training mode, and a random number of quantizers is used.
  Returns
  -------
  z : Tensor[B x D x T]
      Quantized continuous representation of input
  codes : Tensor[B x N x T]
      Codebook indices for each codebook
      (quantized discrete representation of input)
  latents : Tensor[B x N*D x T]
      Projected latents (continuous representation of input before quantization)
  vq/commitment_loss : Tensor[1]
      Commitment loss to train encoder to predict vectors closer to codebook
      entries
  vq/codebook_loss : Tensor[1]
      Codebook loss to update the codebook
docstr_tok: ["Quantized","the","input","tensor","using","a","fixed","set","of","`","n","`","codebooks","and","returns","the","corresponding","codebook","vectors","Parameters","--","--","--","--","--","z",":","Tensor","[","B","x","D","x","T","]","n_quantizers",":","int",",","optional","No",".","of","quantizers","to","use","(","n_quantizers","<","self.n_codebooks","ex",":","for","quantizer","dropout",")","Note",":","if","`","self.quantizer_dropout","`","is","True",",","this","argument","is","ignored","when","in","training","mode",",","and","a","random","number","of","quantizers","is","used",".","Returns","--","--","--","-","z",":","Tensor","[","B","x","D","x","T","]","Quantized","continuous","representation","of","input","codes",":","Tensor","[","B","x","N","x","T","]","Codebook","indices","for","each","codebook","(","quantized","discrete","representation","of","input",")","latents",":","Tensor","[","B","x","N","*","D","x","T","]","Projected","latents","(","continuous","representation","of","input","before","quantization",")","vq\/commitment_loss",":","Tensor","[","1","]","Commitment","loss","to","train","encoder","to","predict","vectors","closer","to","codebook","entries","vq\/codebook_loss",":","Tensor","[","1","]","Codebook","loss","to","update","the","codebook"]
code: def forward(self, z, n_quantizers: int = None):
      """Quantized the input tensor using a fixed set of `n` codebooks and returns
      the corresponding codebook vectors
      Parameters
      ----------
      z : Tensor[B x D x T]
      n_quantizers : int, optional
          No. of quantizers to use
          (n_quantizers < self.n_codebooks ex: for quantizer dropout)
          Note: if `self.quantizer_dropout` is True, this argument is ignored
              when in training mode, and a random number of quantizers is used.
      Returns
      -------
      z : Tensor[B x D x T]
          Quantized continuous representation of input
      codes : Tensor[B x N x T]
          Codebook indices for each codebook
          (quantized discrete representation of input)
      latents : Tensor[B x N*D x T]
          Projected latents (continuous representation of input before quantization)
      vq/commitment_loss : Tensor[1]
          Commitment loss to train encoder to predict vectors closer to codebook
          entries
      vq/codebook_loss : Tensor[1]
          Codebook loss to update the codebook
      """
      z_q = 0
      residual = z
      commitment_loss = 0
      codebook_loss = 0
  
      codebook_indices = []
      latents = []
  
      if n_quantizers is None:
          n_quantizers = self.n_codebooks
      if self.training:
          n_quantizers = (
              torch.ones((z.shape[0],)) * self.n_codebooks + 1
          )
          dropout = torch.randint(
              1, self.n_codebooks + 1, (z.shape[0],)
          )
          n_dropout = int(z.shape[0] * self.quantizer_dropout)
          n_quantizers[:n_dropout] = dropout[:n_dropout]
          n_quantizers = n_quantizers.to(z.device)
  
      for i, quantizer in enumerate(self.quantizers):
          if self.training is False and i >= n_quantizers:
              break
  
          (
              z_q_i,
              commitment_loss_i,
              codebook_loss_i,
              indices_i,
              z_e_i,
          ) = quantizer(residual)
  
          # Create mask to apply quantizer dropout
          mask = (
              torch.full((z.shape[0],), fill_value=i, device=z.device)
              < n_quantizers
          )
          z_q = z_q + z_q_i * mask[:, None, None]
          residual = residual - z_q_i
  
          # Sum losses
          commitment_loss += (commitment_loss_i * mask).mean()
          codebook_loss += (codebook_loss_i * mask).mean()
  
          codebook_indices.append(indices_i)
          latents.append(z_e_i)
  
      codes = torch.stack(codebook_indices, dim=1)
      latents = torch.cat(latents, dim=1)
  
      return z_q, codes, latents, commitment_loss, codebook_loss
  
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 ResidualVectorQuantize obj's underlying member method code with example usage: obj.forward(z, n_quantizers) and returns: z_q,  codes,  latents,  commitment_loss,  codebook_loss