code: def test_scatter(reduce, device):
      import torch_scatter
  
      src = torch.randn(100, 16, device=device)
      index = torch.randint(0, 8, (100,), device=device)
  
      out1 = scatter(src, index, dim=0, reduce=reduce)
      out2 = torch_scatter.scatter(src, index, dim=0, reduce=reduce)
      assert out1.device == device
      assert torch.allclose(out1, out2, atol=1e-6)
  
      jit = torch.jit.script(scatter)
      out3 = jit(src, index, dim=0, reduce=reduce)
      assert torch.allclose(out1, out3, atol=1e-6)
  
      src = torch.randn(8, 100, 16, device=device)
      out1 = scatter(src, index, dim=1, reduce=reduce)
      out2 = torch_scatter.scatter(src, index, dim=1, reduce=reduce)
      assert out1.device == device
      assert torch.allclose(out1, out2, atol=1e-6)
  
file_path: pytorch_geometric/test/utils/test_scatter.py
imports: from itertools import product
  import pytest
  import torch
  import torch_geometric.typing
  from torch_geometric.profile import benchmark
  from torch_geometric.testing import withCUDA, withPackage
  from torch_geometric.utils import group_argsort, scatter
  from torch_geometric.utils._scatter import scatter_argmax
label_desc: Use image node_id 2 for calling a global function with example usage: test_scatter(reduce, device) without return types