class_docstr: Main model class.
class_docstr_tok: ["Main","model","class","."]
docstr: Computes rate and distortion losses.
docstr_tok: ["Computes","rate","and","distortion","losses","."]
code: def call(self, x, training):
      """Computes rate and distortion losses."""
      entropy_model = tfc.ContinuousBatchedEntropyModel(
          self.prior, coding_rank=3, compression=False
      )
      x = tf.cast(
          x, self.compute_dtype
      )  # TODO(jonycgn): Why is this necessary?
      y = self.analysis_transform(x)
      y_hat, bits = entropy_model(y, training=training)
      x_hat = self.synthesis_transform(y_hat)
      # Total number of bits divided by total number of pixels.
      num_pixels = tf.cast(tf.reduce_prod(tf.shape(x)[:-1]), bits.dtype)
      bpp = tf.reduce_sum(bits) / num_pixels
      # Mean squared error across pixels.
      mse = tf.reduce_mean(tf.math.squared_difference(x, x_hat))
      mse = tf.cast(mse, bpp.dtype)
      # The rate-distortion Lagrangian.
      loss = bpp + self.lmbda * mse
      return loss, bpp, mse
  
imports: import argparse
  import glob
  import sys
  from absl import app
  from absl.flags import argparse_flags
  import tensorflow
  import tensorflow_compression
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling the BLS2017Model obj's underlying member method code with example usage: obj.call(x, training) and returns: loss,  bpp,  mse