code: def test_knn(dtype, device):
      x = tensor(
          [
              [-1, -1],
              [-1, +1],
              [+1, +1],
              [+1, -1],
              [-1, -1],
              [-1, +1],
              [+1, +1],
              [+1, -1],
          ],
          dtype,
          device,
      )
      y = tensor(
          [
              [1, 0],
              [-1, 0],
          ],
          dtype,
          device,
      )
  
      batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
      batch_y = tensor([0, 1], torch.long, device)
  
      edge_index = knn(x, y, 2)
      assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
  
      jit = torch.jit.script(knn)
      edge_index = jit(x, y, 2)
      assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
  
      edge_index = knn(x, y, 2, batch_x, batch_y)
      assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
  
      if x.is_cuda:
          edge_index = knn(x, y, 2, batch_x, batch_y, cosine=True)
          assert to_set(edge_index) == set(
              [(0, 2), (0, 3), (1, 4), (1, 5)]
          )
  
      # Skipping a batch
      batch_x = tensor([0, 0, 0, 0, 2, 2, 2, 2], torch.long, device)
      batch_y = tensor([0, 2], torch.long, device)
      edge_index = knn(x, y, 2, batch_x, batch_y)
      assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
  
file_path: pytorch_cluster/test/test_knn.py
imports: from itertools import product
  import pytest
  import scipy.spatial
  import torch
  from torch_cluster import knn, knn_graph
  from torch_cluster.testing import devices, grad_dtypes, tensor
label_desc: Use image node_id 2 for calling a global function with example usage: test_knn(dtype, device) without return types