code: def _build_codebook(
      codebook_size: int,
      vector_length: int,
      init_distribution: str = "uniform",
  ) -> Tensor:
      if init_distribution == "uniform":
          codebook = torch.empty(codebook_size, vector_length).uniform_(
              -1.0 / codebook_size, 1.0 / vector_length
          )
      elif init_distribution == "normal":
          codebook = torch.randn(codebook_size, vector_length)
      else:
          raise ValueError(
              f"Unrecognized init_distribution {init_distribution}."
          )
  
      return codebook
  
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 1 for calling a global function with example usage: _build_codebook(codebook_size, vector_length, init_distribution) and returns: codebook