def _kl_von_mises_von_mises(d1, d2, name=None):
  """Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)
  """
  with tf.name_scope(name or "kl_von_mises_von_mises"):
    # The density of von Mises is (abbreviating the concentration for conc):
    #   vonMises(x; loc, conc) = exp(conc cos(x - loc)) / (2 pi I_0 (conc) )
    # We need two properties:
    # 1. Standardization: if z ~ vonMises(0, conc), then
    #    z + loc ~ vonMises(loc, conc).
    # 2. Expectation of cosine:
    #    E_q(z | 0, conc) cos z = I_1 (conc) / I_0 (conc)
    # Now,
    # KL(d1 || d2)
    #   = E_vonMises(x; loc1, conc1) log vonMises(x; loc1, conc1)
    #                                     / vonMises(x; loc2, conc2)
    # Plugging the densities and rearranging, we have
    #   log I_0(conc2) / I_0(conc1)
    #     + E_vonMises(x; loc1, conc1) [ conc1 cos (z - loc1)
    #                                     - conc2 cos (z - loc2) ]
    # Let's transform the second term using the standardization property:
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (z - (loc2 - loc1))]
    # Applying the cos (x - y) = cos x cos y + sin x sin y expansion, we get
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (loc2 - loc1) cos z
    #     - conc2 sin(loc2 - loc1) sin z]
    # Because the distribution is symmetric around zero, the last term vanishes
    # in expectation. The remaining two terms are computed using the
    # "expectation of cosine" property:
    #   (conc1 - conc2 cos (loc2 - loc1) E_vonMises(x; 0, conc1) cos z
    #     = (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # In total, we have
    #   KL(d1 || d2) = log I_0(conc2) / I_0(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # To improve the numerical stability, we can replace I_j(k) functions with
    # the exponentially scaled versions using the equality
    # I_j(k) = I_j^E(k) exp(k) (which holds for k >= 0):
    #   KL(d1 || d2) = (conc2 - conc1) + log I_0^E(conc2) / I_0^E(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1^E(conc1) / I_0^E(conc1)
    # Note that this formula is numerically stable for conc1 = 0 and/or
    # conc2 = 0 because I_0 (0) = I_0^E (0) = 1.
    i0e_concentration1 = tf.math.bessel_i0e(d1.concentration)
    i1e_concentration1 = tf.math.bessel_i1e(d1.concentration)
    i0e_concentration2 = tf.math.bessel_i0e(d2.concentration)
    return ((d2.concentration - d1.concentration) +
            tf.math.log(i0e_concentration2 / i0e_concentration1) +
            (d1.concentration - d2.concentration * tf.cos(d1.loc - d2.loc)) *
            (i1e_concentration1 / i0e_concentration1))
def _kl_von_mises_von_mises(d1, d2, name=None):
  """Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)
  """
  with tf.name_scope(name or "kl_von_mises_von_mises"):
    # The density of von Mises is (abbreviating the concentration for conc):
    #   vonMises(x; loc, conc) = exp(conc cos(x - loc)) / (2 pi I_0 (conc) )
    # We need two properties:
    # 1. Standardization: if z ~ vonMises(0, conc), then
    #    z + loc ~ vonMises(loc, conc).
    # 2. Expectation of cosine:
    #    E_q(z | 0, conc) cos z = I_1 (conc) / I_0 (conc)
    # Now,
    # KL(d1 || d2)
    #   = E_vonMises(x; loc1, conc1) log vonMises(x; loc1, conc1)
    #                                     / vonMises(x; loc2, conc2)
    # Plugging the densities and rearranging, we have
    #   log I_0(conc2) / I_0(conc1)
    #     + E_vonMises(x; loc1, conc1) [ conc1 cos (z - loc1)
    #                                     - conc2 cos (z - loc2) ]
    # Let's transform the second term using the standardization property:
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (z - (loc2 - loc1))]
    # Applying the cos (x - y) = cos x cos y + sin x sin y expansion, we get
    #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (loc2 - loc1) cos z
    #     - conc2 sin(loc2 - loc1) sin z]
    # Because the distribution is symmetric around zero, the last term vanishes
    # in expectation. The remaining two terms are computed using the
    # "expectation of cosine" property:
    #   (conc1 - conc2 cos (loc2 - loc1) E_vonMises(x; 0, conc1) cos z
    #     = (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # In total, we have
    #   KL(d1 || d2) = log I_0(conc2) / I_0(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1)
    # To improve the numerical stability, we can replace I_j(k) functions with
    # the exponentially scaled versions using the equality
    # I_j(k) = I_j^E(k) exp(k) (which holds for k >= 0):
    #   KL(d1 || d2) = (conc2 - conc1) + log I_0^E(conc2) / I_0^E(conc1)
    #     + (conc1 - conc2 cos (loc2 - loc1)) I_1^E(conc1) / I_0^E(conc1)
    # Note that this formula is numerically stable for conc1 = 0 and/or
    # conc2 = 0 because I_0 (0) = I_0^E (0) = 1.
    i0e_concentration1 = tf.math.bessel_i0e(d1.concentration)
    i1e_concentration1 = tf.math.bessel_i1e(d1.concentration)
    i0e_concentration2 = tf.math.bessel_i0e(d2.concentration)
    return ((d2.concentration - d1.concentration) +
            tf.math.log(i0e_concentration2 / i0e_concentration1) +
            (d1.concentration - d2.concentration * tf.cos(d1.loc - d2.loc)) *
            (i1e_concentration1 / i0e_concentration1))
def _kl_von_mises_von_mises ( d1 , d2 , name = None ) : with tf . name_scope ( name or "kl_von_mises_von_mises" ) : # The density of von Mises is (abbreviating the concentration for conc): #   vonMises(x; loc, conc) = exp(conc cos(x - loc)) / (2 pi I_0 (conc) ) # We need two properties: # 1. Standardization: if z ~ vonMises(0, conc), then #    z + loc ~ vonMises(loc, conc). # 2. Expectation of cosine: #    E_q(z | 0, conc) cos z = I_1 (conc) / I_0 (conc) # Now, # KL(d1 || d2) #   = E_vonMises(x; loc1, conc1) log vonMises(x; loc1, conc1) #                                     / vonMises(x; loc2, conc2) # Plugging the densities and rearranging, we have #   log I_0(conc2) / I_0(conc1) #     + E_vonMises(x; loc1, conc1) [ conc1 cos (z - loc1) #                                     - conc2 cos (z - loc2) ] # Let's transform the second term using the standardization property: #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (z - (loc2 - loc1))] # Applying the cos (x - y) = cos x cos y + sin x sin y expansion, we get #   E_vonMises(x; 0, conc1) [conc1 cos z - conc2 cos (loc2 - loc1) cos z #     - conc2 sin(loc2 - loc1) sin z] # Because the distribution is symmetric around zero, the last term vanishes # in expectation. The remaining two terms are computed using the # "expectation of cosine" property: #   (conc1 - conc2 cos (loc2 - loc1) E_vonMises(x; 0, conc1) cos z #     = (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1) # In total, we have #   KL(d1 || d2) = log I_0(conc2) / I_0(conc1) #     + (conc1 - conc2 cos (loc2 - loc1)) I_1(conc1) / I_0(conc1) # To improve the numerical stability, we can replace I_j(k) functions with # the exponentially scaled versions using the equality # I_j(k) = I_j^E(k) exp(k) (which holds for k >= 0): #   KL(d1 || d2) = (conc2 - conc1) + log I_0^E(conc2) / I_0^E(conc1) #     + (conc1 - conc2 cos (loc2 - loc1)) I_1^E(conc1) / I_0^E(conc1) # Note that this formula is numerically stable for conc1 = 0 and/or # conc2 = 0 because I_0 (0) = I_0^E (0) = 1. i0e_concentration1 = tf . math . bessel_i0e ( d1 . concentration ) i1e_concentration1 = tf . math . bessel_i1e ( d1 . concentration ) i0e_concentration2 = tf . math . bessel_i0e ( d2 . concentration ) return ( ( d2 . concentration - d1 . concentration ) + tf . math . log ( i0e_concentration2 / i0e_concentration1 ) + ( d1 . concentration - d2 . concentration * tf . cos ( d1 . loc - d2 . loc ) ) * ( i1e_concentration1 / i0e_concentration1 ) )
Batchwise KL divergence KL(d1 || d2) with d1 and d2 von Mises.

  Args:
    d1: instance of a von Mises distribution object.
    d2: instance of a a von Mises distribution object.
    name: (optional) Name to use for created operations.
      default is "kl_von_mises_von_mises".

  Returns:
    Batchwise KL(d1 || d2)
Batchwise KL divergence KL ( d1 || d2 ) with d1 and d2 von Mises .