code: def test_segment(device, reduce):
      src = torch.randn(20, 16, device=device)
      ptr = torch.tensor([0, 5, 10, 15, 20], device=device)
  
      out = segment(src, ptr, reduce=reduce)
  
      expected = getattr(torch, reduce)(src.view(4, 5, -1), dim=1)
      expected = (
          expected[0] if isinstance(expected, tuple) else expected
      )
  
      assert torch.allclose(out, expected)
  
imports: import pytest
  import torch
  from torch_geometric.testing import withCUDA, withPackage
  from torch_geometric.utils import segment
label_desc: Use image node_id 1 for calling a global function with example usage: test_segment(device, reduce) without return types