code: def gmm_energy(
      z: tf.Tensor,
      phi: tf.Tensor,
      mu: tf.Tensor,
      cov: tf.Tensor,
      L: tf.Tensor,
      log_det_cov: tf.Tensor,
      return_mean: bool = True,
  ) -> Tuple[tf.Tensor, tf.Tensor]:
      """
      Compute sample energy from Gaussian Mixture Model.
  
      Parameters
      ----------
      z
          Observations.
      phi
          Mixture component distribution weights.
      mu
          Mixture means.
      cov
          Mixture covariance.
      L
          Cholesky decomposition of `cov`.
      log_det_cov
          Log of the determinant of `cov`.
      return_mean
          Take mean across all sample energies in a batch.
  
      Returns
      -------
      sample_energy
          The sample energy of the GMM.
      cov_diag
          The inverse sum of the diagonal components of the covariance matrix.
      """
      D = tf.shape(cov)[1]
      z_mu = tf.expand_dims(z, 1) - tf.expand_dims(mu, 0)  # N x K x D
      z_mu_T = tf.transpose(z_mu, perm=[1, 2, 0])  # K x D x N
      v = tf.linalg.triangular_solve(L, z_mu_T, lower=True)  # K x D x D
  
      # rewrite sample energy in logsumexp format for numerical stability
      logits = tf.math.log(tf.expand_dims(phi, -1)) - 0.5 * (
          tf.reduce_sum(tf.square(v), 1)
          + tf.cast(D, tf.float32) * tf.math.log(2.0 * np.pi)
          + tf.expand_dims(log_det_cov, -1)
      )  # K x N
      sample_energy = -tf.reduce_logsumexp(logits, axis=0)  # N
  
      if return_mean:
          sample_energy = tf.reduce_mean(sample_energy)
  
      # inverse sum of variances
      cov_diag = tf.reduce_sum(tf.divide(1, tf.linalg.diag_part(cov)))
  
      return sample_energy, cov_diag
  
imports: import numpy
  import tensorflow
  from typing import Tuple
label_desc: Use image node_id 2 for calling a global function with example usage: gmm_energy(z, phi, mu, cov, L, log_det_cov, return_mean) and returns: sample_energy,  cov_diag