code: def compute_nats_and_bits_per_dim(
      data_dim, latent_dim, average_reconstruction, average_prior
  ):
      """Computes negative ELBO, which is an upper bound on the negative likelihood.
  
      Args:
        data_dim: int-like indicating data dimensionality.
        latent_dim: int-like indicating latent dimensionality.
        average_reconstruction: Scalar Tensor indicating the reconstruction cost
          averaged over all data dimensions and any data batches.
        average_prior: Scalar Tensor indicating the negative log-prior probability
          averaged over all latent dimensions and any data batches.
  
      Returns:
        Tuple of scalar Tensors, representing the nats and bits per data dimension
        (e.g., subpixels) respectively.
      """
      with tf.name_scope(None, default_name="compute_nats_per_dim"):
          data_dim = tf.cast(data_dim, average_reconstruction.dtype)
          latent_dim = tf.cast(latent_dim, average_prior.dtype)
          negative_log_likelihood = data_dim * average_reconstruction
          negative_log_prior = latent_dim * average_prior
          negative_elbo = negative_log_likelihood + negative_log_prior
          nats_per_dim = tf.divide(
              negative_elbo, data_dim, name="nats_per_dim"
          )
          bits_per_dim = tf.divide(
              nats_per_dim, tf.log(2.0), name="bits_per_dim"
          )
          return nats_per_dim, bits_per_dim
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_image_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import transformer_layers
  from tensor2tensor.utils import beam_search
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  import tensorflow_probability
label_desc: Use image node_id 2 for calling a global function with example usage: compute_nats_and_bits_per_dim(data_dim, latent_dim, average_reconstruction, average_prior) and returns: nats_per_dim,  bits_per_dim