code: def pairwise_kl_loss(
      logits: Array,
      neighbourhood_logits: Array,
      knn_indices: Array,
      knn_similarities: Array,
      temperature: float,
      epsilon: float = 1e-16,
      example_weights: Optional[Array] = None,
  ) -> float:
      """KL Divergence loss weighted by similarity within a neighbourhood.
  
      Args:
        logits: [n, d] array.
        neighbourhood_logits: [m, d] array.
        knn_indices: [n, k] array of nearest neighbours for each logit. Indexes
          into "neighbourhood_logits". Therefore, in each row, each of the k values
          are integers in the interval [0, m).
        knn_similarities: [n, k] array of each of the similarity scores
          corresponding to the knn_indices.
        temperature: A float.
        epsilon: Small constant for numerical stability.
        example_weights: The weight for each example in the batch. Can be used
          to account for TPU padding.
  
      Returns:
        The KL divergence of each logit to its neighbourhood, weighted by the
          similarity.
      """
      n, d = logits.shape
      k = knn_indices.shape[1]
  
      knn_logits = neighbourhood_logits[
          knn_indices.reshape(-1)
      ].reshape([n, k, d])
  
      t_softmax_temp = (
          jax.nn.softmax(knn_logits / temperature) + epsilon
      )
      s_softmax_temp = jax.nn.log_softmax(
          logits / temperature
      )  # [n, d]
  
      # Normalize the sum of similarities by their sum, so that the new labels
      # will sum up to 1.
      normalized_sim = knn_similarities / jnp.reshape(
          jnp.sum(knn_similarities, axis=-1), (-1, 1)
      )  # [n, k]
  
      # Multiply the labels by their corresponding similarity value.
      weighted_t_softmax = jnp.squeeze(
          jnp.matmul(jnp.expand_dims(normalized_sim, 1), t_softmax_temp)
      )  # [n, d]
  
      kldiv_loss_per_pair = weighted_t_softmax * (
          jnp.log(weighted_t_softmax) - s_softmax_temp
      )  # [n, m]
      kldiv_loss_per_example = jnp.power(temperature, 2) * jnp.sum(
          kldiv_loss_per_pair, 1
      )  # [n, 1]
  
      if example_weights is not None:
          normalization = example_weights.sum()
      else:
          normalization = n
      return jnp.sum(kldiv_loss_per_example) / (normalization + epsilon)
  
imports: from typing import Optional, Tuple, Union
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: pairwise_kl_loss(logits, neighbourhood_logits, knn_indices, knn_similarities, temperature, epsilon, example_weights) and returns: unknown