class_docstr: Nonlinear transform coding model.
code: def quantize(self, x, **kwargs):
      y_hat, x_hat, rates = self.encode_decode(
          x, False, False, False, **kwargs
      )
      # Find the unique set of latents for these inputs. Converts integer indexes
      # on the infinite lattice to scalar indexes into a codebook (which is only
      # valid for this set of inputs).
      _, i, indexes = np.unique(
          tf.reshape(y_hat, [-1, self.ndim_latent]),
          return_index=True,
          return_inverse=True,
          axis=0,
      )
      codebook = tf.gather(
          tf.reshape(x_hat, [-1, self.ndim_source]), i, axis=0
      )
      rates = tf.gather(tf.reshape(rates, [-1]), i)
      indexes = tf.reshape(tf.cast(indexes, tf.int32), tf.shape(x)[:-1])
      return codebook, rates, indexes
  
imports: import matplotlib.pyplot
  import numpy
  import tensorflow
  import tensorflow_compression
  import tensorflow_probability
  from toy_sources import compression_model
label_desc: Use image node_id 11 for calling the NTCModel obj's underlying member method code with example usage: obj.quantize(x) and returns: codebook,  rates,  indexes