code: def _kl_brute_force(a, b, name=None):
      """Batched KL divergence `KL(a || b)` for multivariate Normals.
  
      With `X`, `Y` both multivariate Normals in `R^k` with means `mu_a`, `mu_b` and
      covariance `C_a`, `C_b` respectively,
  
      ```
      KL(a || b) = 0.5 * ( L - k + T + Q ),
      L := Log[Det(C_b)] - Log[Det(C_a)]
      T := trace(C_b^{-1} C_a),
      Q := (mu_b - mu_a)^T C_b^{-1} (mu_b - mu_a),
      ```
  
      This `Op` computes the trace by solving `C_b^{-1} C_a`. Although efficient
      methods for solving systems with `C_b` may be available, a dense version of
      (the square root of) `C_a` is used, so performance is `O(B s k**2)` where `B`
      is the batch size, and `s` is the cost of solving `C_b x = y` for vectors `x`
      and `y`.
  
      Args:
        a: Instance of `MultivariateNormalLinearOperator`.
        b: Instance of `MultivariateNormalLinearOperator`.
        name: (optional) name to use for created ops. Default "kl_mvn".
  
      Returns:
        Batchwise `KL(a || b)`.
      """
  
      def squared_frobenius_norm(x):
          """Helper to make KL calculation slightly more readable."""
          # http://mathworld.wolfram.com/FrobeniusNorm.html
          # The gradient of KL[p,q] is not defined when p==q. The culprit is
          # tf.norm, i.e., we cannot use the commented out code.
          # return tf.square(tf.norm(x, ord="fro", axis=[-2, -1]))
          return tf.reduce_sum(tf.square(x), axis=[-2, -1])
  
      # TODO(b/35041439): See also b/35040945. Remove this function once LinOp
      # supports something like:
      #   A.inverse().solve(B).norm(order='fro', axis=[-1, -2])
      def is_diagonal(x):
          """Helper to identify if `LinearOperator` has only a diagonal component."""
          return (
              isinstance(x, tf.linalg.LinearOperatorIdentity)
              or isinstance(x, tf.linalg.LinearOperatorScaledIdentity)
              or isinstance(x, tf.linalg.LinearOperatorDiag)
          )
  
      with tf.name_scope(name or "kl_mvn"):
          # Calculation is based on:
          # http://stats.stackexchange.com/questions/60680/kl-divergence-between-two-multivariate-gaussians
          # and,
          # https://en.wikipedia.org/wiki/Matrix_norm#Frobenius_norm
          # i.e.,
          #   If Ca = AA', Cb = BB', then
          #   tr[inv(Cb) Ca] = tr[inv(B)' inv(B) A A']
          #                  = tr[inv(B) A A' inv(B)']
          #                  = tr[(inv(B) A) (inv(B) A)']
          #                  = sum_{ij} (inv(B) A)_{ij}**2
          #                  = ||inv(B) A||_F**2
          # where ||.||_F is the Frobenius norm and the second equality follows from
          # the cyclic permutation property.
          if is_diagonal(a.scale) and is_diagonal(b.scale):
              # Using `stddev` because it handles expansion of Identity cases.
              b_inv_a = (a.stddev() / b.stddev())[..., tf.newaxis]
          else:
              b_inv_a = b.scale.solve(a.scale.to_dense())
          kl_div = (
              b.scale.log_abs_determinant()
              - a.scale.log_abs_determinant()
              + 0.5
              * (
                  -tf.cast(a.scale.domain_dimension_tensor(), a.dtype)
                  + squared_frobenius_norm(b_inv_a)
                  + squared_frobenius_norm(
                      b.scale.solve(
                          (b.mean() - a.mean())[..., tf.newaxis]
                      )
                  )
              )
          )
          tensorshape_util.set_shape(
              kl_div,
              tf.broadcast_static_shape(a.batch_shape, b.batch_shape),
          )
          return kl_div
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import scale_matvec_linear_operator
  from tensorflow_probability.python.bijectors import shift
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: _kl_brute_force(a, b, name) and returns: tf, unknown, kl_div