code: def ncr_loss(
      logits: Array,
      features: Array,
      batch_logits: Array,
      batch_features: Array,
      number_neighbours: int,
      smoothing_gamma: float,
      temperature: float = 1.0,
      example_weights: Optional[Array] = None,
  ) -> float:
      """Computes the Neighbourhood Consistency Regularisation loss.
  
      Details are in: https://arxiv.org/pdf/2202.02200.pdf
  
      Args:
        logits: An [n_batch, n_classes] array.
        features: An [n_batch, d] array.
        batch_logits: An [m, n_classes] array.
        batch_features: An [m, d] array.
        number_neighbours: The number of neighbours to use. Must be an integer in
          the interval [1, m]
        smoothing_gamma: A value in (0, infinity)
        temperature: Temperature for the KL-Divergence. A value in (0, infinity).
        example_weights: If not None, the weight to apply to the loss of each
          example in the batch. Useful for dealing with TPU padding. If None, the
          mean over the batch is computed.
  
      Returns:
        The weighted NCR loss computed over the batch.
      """
  
      indices, similarities = get_knn(
          features, batch_features, number_neighbours + 1
      )
      # Remove the example itself from the list of nearest neighbours.
      indices = indices[:, 1:]
      similarities = similarities[:, 1:]
  
      similarities = jnp.power(similarities, smoothing_gamma)
      loss = pairwise_kl_loss(
          logits,
          batch_logits,
          indices,
          similarities,
          temperature,
          example_weights=example_weights,
      )
      return loss
  
imports: from typing import Optional, Tuple, Union
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: ncr_loss(logits, features, batch_logits, batch_features, number_neighbours, smoothing_gamma, temperature, example_weights) and returns: loss