class_docstr: 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.LocationScaleIndexedEntropyModel(
          tfc.NoisyNormal,
          self.num_scales,
          self.scale_fn,
          coding_rank=3,
          compression=False,
      )
      side_entropy_model = tfc.ContinuousBatchedEntropyModel(
          self.hyperprior, coding_rank=3, compression=False
      )
  
      x = tf.cast(
          x, self.compute_dtype
      )  # TODO(jonycgn): Why is this necessary?
      y = self.analysis_transform(x)
      z = self.hyper_analysis_transform(abs(y))
      z_hat, side_bits = side_entropy_model(z, training=training)
      indexes = self.hyper_synthesis_transform(z_hat)
      y_hat, bits = entropy_model(y, indexes, 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) + tf.reduce_sum(side_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 BMSHJ2018Model obj's underlying member method code with example usage: obj.call(x, training) and returns: loss,  bpp,  mse