code: def test_small_k_query(self):
      num_neighbors = 2
      expected_neighbors = np.array(
          [[[0.0, 0.0], [7.5, 1.0]], [[7.5, 1.0], [0.0, 0.0]]]
      )
      expected_neg_distances = np.array(
          [[-5.69, -30.34], [-1.0, -56.25]]
      )
      expected_neighbor_indices = np.array([[0.0, 1.0], [1.0, 0.0]])
  
      @self.variant
      def query_variant(data, points):
          return episodic_memory.knn_query(data, points, num_neighbors)
  
      actual = query_variant(self.data, self.query_points)
  
      np.testing.assert_allclose(
          actual.neighbors, expected_neighbors, atol=1e-6
      )
      np.testing.assert_allclose(
          actual.neighbor_indices, expected_neighbor_indices, atol=1e-6
      )
      np.testing.assert_allclose(
          actual.neighbor_neg_distances,
          expected_neg_distances,
          atol=1e-6,
      )
  
file_path: rlax/rlax/_src/episodic_memory_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import jax
  import numpy
  from rlax._src import episodic_memory
label_desc: Use image node_id 2 for calling the KNNQueryTest obj's underlying member method code with example usage: obj.test_small_k_query() and returns: episodic_memory