code: def get_knn(
      queries: Array,
      dataset: Array,
      k: int,
      zero_negative_similarities: bool = True,
  ) -> Tuple[Array, Array]:
      """Return k nearest neighbours from dataset given queries.
  
      Args:
        queries: An [q, d] array where q is the number of queries, and d the
          dimensionality of each query-vector.
        dataset: An [n, d] array where n is the number of examples.
        k: The number of nearest neighbours to retrieve.
        zero_negative_similarities: If true, negative similarities are set to 0.
  
      Returns:
        indices: A [q, k] dimensional array. For each query, the k indices of the
          nearest neighbours are returned.
        similarities: A [q, k] dimensional array with the similarities to each
          query. Similarities and corresponding indices are sorted, in descending
          order.
      """
      if k <= 0:
          k = dataset.shape[0]
      if k > dataset.shape[0]:
          k = dataset.shape[0]
  
      queries = l2_normalize(queries, axis=-1)
      dataset = l2_normalize(dataset, axis=-1)
  
      all_similarities = jnp.matmul(
          queries, jnp.transpose(dataset)
      )  # [q, n]
      similarities, indices = jax.lax.top_k(all_similarities, k)
      if zero_negative_similarities:
          similarities = jax.nn.relu(similarities)
  
      return indices, similarities
  
imports: from typing import Optional, Tuple, Union
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: get_knn(queries, dataset, k, zero_negative_similarities) and returns: indices,  similarities