code: def _kl_power_spherical_vmf(a, b, name=None):
      """Calculate the batched KL divergence KL(a || b).
  
      Args:
        a: instance of a PowerSpherical distribution object.
        b: instance of a VonMisesFisher distribution object.
        name: (optional) Name to use for created operations.
          default is "kl_power_spherical_vmf".
  
      Returns:
        Batchwise KL(a || b)
  
      Raises:
        ValueError: If the two distributions are over spheres of different
          dimensions.
  
      #### References
  
      [1] Nicola de Cao, Wilker Aziz. The Power Spherical distribution.
          https://arxiv.org/abs/2006.04437.
      """
      with tf.name_scope(name or "kl_power_spherical_vmf"):
          msg = (
              "Can not compute the KL divergence between a `PowerSpherical` and "
              "`VonMisesFisher` of different dimensions."
          )
          deps = []
          if (
              a.event_shape[-1] is not None
              and b.event_shape[-1] is not None
          ):
              if a.event_shape[-1] != b.event_shape[-1]:
                  raise ValueError(
                      (msg + "Got {} vs. {}").format(
                          a.event_shape[-1], b.event_shape[-1]
                      )
                  )
          elif a.validate_args or b.validate_args:
              deps += [
                  assert_util.assert_equal(
                      a.event_shape_tensor()[-1],
                      b.event_shape_tensor()[-1],
                      message=msg,
                  )
              ]
  
          with tf.control_dependencies(deps):
              a_mean_direction = tf.convert_to_tensor(a.mean_direction)
              a_concentration = tf.convert_to_tensor(a.concentration)
              b_mean_direction = tf.convert_to_tensor(b.mean_direction)
              b_concentration = tf.convert_to_tensor(b.concentration)
  
              event_size = tf.cast(
                  a._event_shape_tensor(  # pylint:disable=protected-access
                      mean_direction=a_mean_direction
                  )[
                      -1
                  ],
                  a.dtype,
              )
              kl = (
                  -a._entropy(
                      concentration=a_concentration,  # pylint:disable=protected-access
                      mean_direction=a_mean_direction,
                  )
                  + b._log_normalization(  # pylint:disable=protected-access
                      concentration=b_concentration
                  )
                  - a_concentration
                  * b_concentration
                  * tf.reduce_sum(
                      a_mean_direction * b_mean_direction, axis=-1
                  )
                  / (a_concentration + event_size - 1.0)
              )
              return kl
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import softmax_centered
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import beta
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import spherical_uniform
  from tensorflow_probability.python.distributions import von_mises_fisher
  from tensorflow_probability.python.internal import assert_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 reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import special
  from tensorflow_probability.python.random import random_ops
label_desc: Use image node_id 2 for calling a global function with example usage: _kl_power_spherical_vmf(a, b, name) and returns: kl