class_docstr: Main model class.
class_docstr_tok: ["Main","model","class","."]
docstr: Compresses an image.
docstr_tok: ["Compresses","an","image","."]
code: def compress(self, x):
      """Compresses an image."""
      # Add batch dimension and cast to float.
      x = tf.expand_dims(x, 0)
      x = tf.cast(x, dtype=self.compute_dtype)
  
      y_strings = []
      x_shape = tf.shape(x)[1:-1]
  
      # Build the encoder (analysis) half of the hierarchical autoencoder.
      y = self.analysis_transform(x)
      y_shape = tf.shape(y)[1:-1]
  
      z = self.hyper_analysis_transform(y)
      z_shape = tf.shape(z)[1:-1]
  
      z_string = self.em_z.compress(z)
      z_hat = self.em_z.decompress(z_string, z_shape)
  
      # Build the decoder (synthesis) half of the hierarchical autoencoder.
      latent_scales = self.hyper_synthesis_scale_transform(z_hat)
      latent_means = self.hyper_synthesis_mean_transform(z_hat)
  
      # En/Decode each slice conditioned on hyperprior and previous slices.
      y_slices = tf.split(y, self.num_slices, axis=-1)
      y_hat_slices = []
      for slice_index, y_slice in enumerate(y_slices):
          # Model may condition on only a subset of previous slices.
          support_slices = (
              y_hat_slices
              if self.max_support_slices < 0
              else y_hat_slices[: self.max_support_slices]
          )
  
          # Predict mu and sigma for the current slice.
          mean_support = tf.concat(
              [latent_means] + support_slices, axis=-1
          )
          mu = self.cc_mean_transforms[slice_index](mean_support)
          mu = mu[:, : y_shape[0], : y_shape[1], :]
  
          # Note that in this implementation, `sigma` represents scale indices,
          # not actual scale values.
          scale_support = tf.concat(
              [latent_scales] + support_slices, axis=-1
          )
          sigma = self.cc_scale_transforms[slice_index](scale_support)
          sigma = sigma[:, : y_shape[0], : y_shape[1], :]
  
          slice_string = self.em_y.compress(y_slice, sigma, mu)
          y_strings.append(slice_string)
          y_hat_slice = self.em_y.decompress(slice_string, sigma, mu)
  
          # Add latent residual prediction (LRP).
          lrp_support = tf.concat([mean_support, y_hat_slice], axis=-1)
          lrp = self.lrp_transforms[slice_index](lrp_support)
          lrp = 0.5 * tf.math.tanh(lrp)
          y_hat_slice += lrp
  
          y_hat_slices.append(y_hat_slice)
  
      return (x_shape, y_shape, z_shape, z_string) + tuple(y_strings)
  
imports: import argparse
  import functools
  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 8 for calling the MS2020Model obj's underlying member method code with example usage: obj.compress(x) and returns: unknown