code: def contrastive_loss(
      query_emb: jnp.ndarray,
      key_emb: jnp.ndarray,
      temperature: float = 1.0,
  ):
      """Contrastive loss with hard negative samples & other in-batch negatives.
  
      Args:
        query_emb: An array of shape [bsz, n_dim].
        key_emb: An array of shape [bsz, n_knowledge, n_dim]. Only the first one is
          true positive sample, and the others are hard negatives.
        temperature: A scalar that the temprature is divided by it.
  
      Returns:
        Computed loss value.
      """
      if query_emb.shape[0] != key_emb.shape[0]:
          raise ValueError(
              "query_emb and key_emb should have the same batch size."
          )
      if query_emb.shape[-1] != key_emb.shape[-1]:
          raise ValueError(
              "query_emb and key_emb should have the same embedding size."
          )
      per_device_bsz, k = query_emb.shape[0], key_emb.shape[1]
      global_key_emb = jnp.concatenate(
          jax.lax.all_gather(key_emb, "batch"), 0
      )
      labels = jax.lax.axis_index(
          axis_name="batch"
      ) * per_device_bsz * k + jnp.arange(per_device_bsz)
  
      # bsz×d @ (bsz*n_device)×K×d -> bsz×(bsz * k * n_device)
      # positive pairs are on first diagonal.
      score_matrix = jnp.reshape(
          jnp.einsum("bd,nkd->bkn", query_emb, global_key_emb),
          [per_device_bsz, -1],
      )
  
      loss = nll_loss(pred=score_matrix / temperature, targets=labels)
      accs = jnp.equal(jnp.argmax(score_matrix, axis=1), labels)
      s0, s1 = score_matrix[0][0], score_matrix[0][1]  # debug purpose
      logging.info("backward host_id : %d", jax.process_index())
      logging.info(jax.lax.axis_index(axis_name="batch"))
      return loss, (jnp.mean(accs), s0, s1)
  
imports: from absl import logging
  from flax.training import common_utils
  import jax
  import jax.numpy
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 2 for calling a global function with example usage: contrastive_loss(query_emb, key_emb, temperature) and returns: loss,  jnp,  s0,  s1