code: def knn_query(
      data: Array,
      query_points: Array,
      num_neighbors: int,
      metric: MetricFn = _sqeuclidian,
  ) -> KNNQueryResult:
      """Finds closest neighbors in data to the query points & their neg distances.
  
      NOTE: For this function to be jittable, static_argnums=[2,] must be passed, as
      the internal jax.lax.top_k(neg_distances, num_neighbors) computation cannot be
      jitted with a dynamic num_neighbors that is passed as an argument.
  
      Args:
        data: array of existing data points (elements in database x feature size)
        query_points: array of points to find neighbors of
            (num query points x feature size).
        num_neighbors: number of neighbors to find.
        metric: Metric to use in calculating distance between two points.
  
      Returns:
        KNNQueryResult with (all sorted by neg distance):
          - neighbors (num query points x num neighbors x feature size)
          - neighbor_indices (num query points x num neighbors)
          - neighbor_neg_distances (num query points x num neighbors)
       * if num_neighbors is greater than number of elements in the database,
         just return KNNQueryResult with the number of elements in the database.
      """
      chex.assert_rank([data, query_points], 2)
      assert data.shape[-1] == query_points.shape[-1]
      distance_fn = jax.jit(functools.partial(_cdist, metric=metric))
      neg_distances = -distance_fn(query_points, data)
      neg_distances, indices = jax.lax.top_k(
          neg_distances, k=min(num_neighbors, data.shape[0])
      )
      # Batch index into data using indices shaped [num queries, num neighbors]
      neighbors = jax.vmap(lambda d, i: d[i], (None, 0))(
          jnp.array(data), indices
      )
      return KNNQueryResult(
          neighbors=neighbors,
          neighbor_indices=indices,
          neighbor_neg_distances=neg_distances,
      )
  
file_path: rlax/rlax/_src/episodic_memory.py
imports: import functools
  from typing import Callable
  import chex
  import jax
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: knn_query(data, query_points, num_neighbors, metric) and returns: KNNQueryResult