code: def _vectorized_kl_divergence(
      h: GaussianDistribution,
      r: GaussianDistribution,
      t: GaussianDistribution,
      exact: bool = True,
  ) -> torch.FloatTensor:
      r"""Vectorized implementation of KL-divergence.
  
      Computes the divergence between :math:`\mathcal{N}(\mu_e, \Sigma_e)` and :math:`\mathcal{N}(\mu_r, \Sigma_r)`
      given by
  
      .. math ::
          \mu_e = \mu_h - \mu_t
  
          \Sigma_e = \Sigma_h + \Sigma_t
  
      where all covariance matrices are diagonal. Hence we can simplify
  
      .. math ::
          D(\mathcal{N}(\mu_e, \Sigma_e), \mathcal{N}(\mu_r, \Sigma_r))
          =
          0.5 * (
            \trace(\Sigma_r^-1 \Sigma_e)
            + (\mu_r - \mu_e) * \Sigma_r^-1 (\mu_r - \mu_e)
            - k
            + \ln (\det(\Sigma_r) / \det(\Sigma_e))
          )
          = 0.5 * (
            \sum_i \Sigma_e[i] / Sigma_r[i]
            + \sum_i \mu[i]^2 / \Sigma_r[i]
            + \sum_i \ln Sigma_r[i]
            - \sum_i \ln Sigma_e[i]
            - k
          )
  
      where :math:`\mu = \mu_r - \mu_e = \mu_r - \mu_h + \mu_t`
  
      :param h: shape: (batch_size, num_heads, 1, 1, d)
          The head entity Gaussian distribution.
      :param r: shape: (batch_size, 1, num_relations, 1, d)
          The relation Gaussian distribution.
      :param t: shape: (batch_size, 1, 1, num_tails, d)
          The tail entity Gaussian distribution.
      :param exact:
          Whether to return the exact similarity, or leave out constant offsets.
  
      :return: torch.Tensor, shape: (s_1, ..., s_k)
          The KL-divergence.
      """
      e_var = h.diagonal_covariance + t.diagonal_covariance
      r_var_safe = at_least_eps(r.diagonal_covariance)
      terms = []
      # 1. Component
      # \sum_i \Sigma_e[i] / Sigma_r[i]
      r_var_safe_reciprocal = r_var_safe.reciprocal()
      terms.append(batched_dot(e_var, r_var_safe_reciprocal))
      # 2. Component
      # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0)
      # with mu = (mu_1 - mu_0)
      # = mu * Sigma_1^-1 mu
      # since Sigma_1 is diagonal
      # = mu**2 / sigma_1
      mu = tensor_sum(r.mean, -h.mean, t.mean)
      terms.append(batched_dot(mu.pow(2), r_var_safe_reciprocal))
      # 3. Component
      if exact:
          terms.append(
              -torch.as_tensor(
                  data=[h.mean.shape[-1]], device=mu.device
              ).squeeze()
          )
      # 4. Component
      # ln (det(\Sigma_1) / det(\Sigma_0))
      # = ln det Sigma_1 - ln det Sigma_0
      # since Sigma is diagonal, we have det Sigma = prod Sigma[ii]
      # = ln prod Sigma_1[ii] - ln prod Sigma_0[ii]
      # = sum ln Sigma_1[ii] - sum ln Sigma_0[ii]
      e_var_safe = at_least_eps(e_var)
      terms.extend(
          (
              r_var_safe.log().sum(dim=-1),
              -e_var_safe.log().sum(dim=-1),
          )
      )
      result = tensor_sum(*terms)
      if exact:
          result = 0.5 * result
      return result
  
file_path: pykeen/src/pykeen/nn/sim.py
imports: import itertools
  import math
  import torch
  from .compute_kernel import batched_dot
  from ..typing import GaussianDistribution
  from ..utils import at_least_eps, calculate_broadcasted_elementwise_result_shape, tensor_sum
label_desc: Use image node_id 3 for calling a global function with example usage: _vectorized_kl_divergence(h, r, t, exact) and returns: result