code: def test_rw_small(device):
      row = tensor([0, 1], torch.long, device)
      col = tensor([1, 0], torch.long, device)
      start = tensor([0, 1, 2], torch.long, device)
      walk_length = 4
  
      out = random_walk(row, col, start, walk_length, num_nodes=3)
      assert out.tolist() == [
          [0, 1, 0, 1, 0],
          [1, 0, 1, 0, 1],
          [2, 2, 2, 2, 2],
      ]
  
      jit = torch.jit.script(random_walk)
      assert torch.equal(
          jit(row, col, start, walk_length, num_nodes=3), out
      )
  
imports: import pytest
  import torch
  from torch_cluster import random_walk
  from torch_cluster.testing import devices, tensor
label_desc: Use image node_id 2 for calling a global function with example usage: test_rw_small(device) without return types