code: def test_radius(dtype, device):
      x = tensor(
          [
              [-1, -1],
              [-1, +1],
              [+1, +1],
              [+1, -1],
              [-1, -1],
              [-1, +1],
              [+1, +1],
              [+1, -1],
          ],
          dtype,
          device,
      )
      y = tensor(
          [
              [0, 0],
              [0, 1],
          ],
          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 = radius(x, y, 2, max_num_neighbors=4)
      assert to_set(edge_index) == set(
          [
              (0, 0),
              (0, 1),
              (0, 2),
              (0, 3),
              (1, 1),
              (1, 2),
              (1, 5),
              (1, 6),
          ]
      )
  
      jit = torch.jit.script(radius)
      edge_index = jit(x, y, 2, max_num_neighbors=4)
      assert to_set(edge_index) == set(
          [
              (0, 0),
              (0, 1),
              (0, 2),
              (0, 3),
              (1, 1),
              (1, 2),
              (1, 5),
              (1, 6),
          ]
      )
  
      edge_index = radius(
          x, y, 2, batch_x, batch_y, max_num_neighbors=4
      )
      assert to_set(edge_index) == set(
          [(0, 0), (0, 1), (0, 2), (0, 3), (1, 5), (1, 6)]
      )
  
      # 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 = radius(
          x, y, 2, batch_x, batch_y, max_num_neighbors=4
      )
      assert to_set(edge_index) == set(
          [(0, 0), (0, 1), (0, 2), (0, 3), (1, 5), (1, 6)]
      )
  
file_path: pytorch_cluster/test/test_radius.py
imports: from itertools import product
  import pytest
  import scipy.spatial
  import torch
  from torch_cluster import radius, radius_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_radius(dtype, device) without return types