class_docstr: Quantization bottleneck.
  
  This class quantizes its input based on a codebook. Each entry is quantized
  to its nearest codebook representation according to Euclidean distance.
  
  Args:
      codebook_sizes: Size of the codebook for each quantization level.
      hidden_dims: Hidden dimension size, i.e., the sizes of the input
          vectors.
      init_distribution: Initialization distribution for the codebooks.
      normalize_codebook: Whether to renormalize the codebook prior to each
          quantization run.
code: def quantize_vectors(
      self,
      preq_latents: Tensor,
      codebook: Tensor,
      probabilistic: bool = False,
  ) -> VQOutput:
      batch_size, hidden_dim, num_y, num_x = preq_latents.shape
      latents = preq_latents.permute(0, 2, 3, 1).reshape(-1, hidden_dim)
  
      # vector is (num_vectors, vector_length), already normalized
      # codebook is (codebook_size, vector_length)
      if self.normalize_codebook:
          codebook = F.normalize(codebook, dim=1)
  
      # distances variable is (num_vectors, codebook_size)
      distances = torch.sum(
          (codebook.unsqueeze(0) - latents.unsqueeze(1)) ** 2, dim=2
      )
  
      if probabilistic:
          # create probability weighted by distance
          weights = 1 / (distances + 1e-7)
          cumulative = torch.cumsum(
              weights / torch.sum(weights, dim=1, keepdim=True), dim=1
          )
          # hacky way to find indices
          # pick approximate indices based on probabilities
          tmp = cumulative - torch.rand_like(cumulative[:, :1])
          tmp[tmp < 0] = 1.0
          indices = torch.argmin(tmp, dim=1)
      else:
          # return the shortest distance
          indices = torch.argmin(distances, dim=1)
  
      # pick out and reshape the latents
      latents = (
          codebook[indices]
          .view(batch_size, num_y, num_x, hidden_dim)
          .permute(0, 3, 1, 2)
          .contiguous()
      )
  
      # VQ losses
      commitment_loss = F.mse_loss(latents.detach(), preq_latents)
      embedding_loss = F.mse_loss(latents, preq_latents.detach())
  
      return VQOutput(
          latent=preq_latents
          + (latents - preq_latents).detach(),  # STE
          prequantized_latent=preq_latents,
          commitment_loss=commitment_loss,
          embedding_loss=embedding_loss,
          indices=indices,
          distances=distances,
      )
  
file_path: NeuralCompression/neuralcompression/layers/_q_bottleneck.py
imports: from typing import NamedTuple, Optional, Sequence
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
label_desc: Use image node_id 4 for calling the QBottleneck obj's underlying member method code with example usage: obj.quantize_vectors(preq_latents, codebook, probabilistic) and returns: VQOutput