class_docstr: Variational entropy-constrained vector quantization model.
code: def __init__(
      self,
      codebook_size,
      initialize="sample",
      logit_scale=1e0,
      **kwargs,
  ):
      super().__init__(**kwargs)
      self.codebook_size = int(codebook_size)
      self.logit_scale = float(logit_scale)
  
      # Initialize codebook with samples from the source distribution.
      if initialize.startswith("sample"):
          codebook_init = self.source.sample(self.codebook_size)
          if len(initialize) > 6:
              scale = float(initialize[7:])
              codebook_init += tf.random.normal(
                  codebook_init.shape, stddev=scale
              )
      elif initialize.startswith("uniform-"):
          width = float(initialize[8:])
          codebook_init = tf.random.uniform(
              [self.codebook_size, self.ndim_source],
              minval=-width / 2,
              maxval=width / 2,
          )
      assert codebook_init.shape == (
          self.codebook_size,
          self.ndim_source,
      )
      logits_init = tf.random.normal(
          [self.codebook_size], stddev=self.logit_scale / 10
      )
      self.codebook = tf.Variable(
          tf.cast(codebook_init, self.dtype), name="codebook"
      )
      self._logits = tf.Variable(
          tf.cast(logits_init, self.dtype), name="logits"
      )
  
file_path: compression/models/toy_sources/vecvq.py
imports: import tensorflow
  from toy_sources import compression_model
label_desc: Use image node_id 1 to create a new VECVQModel object from inherited base classes: compression_model with example: obj = VECVQModel(codebook_size, initialize, logit_scale)