code: def test_forward(test, reduce, dtype, device):
      src = tensor(test["src"], dtype, device)
      index = tensor(test["index"], torch.long, device)
      indptr = tensor(test["indptr"], torch.long, device)
      expected = tensor(test[reduce], dtype, device)
  
      fn = getattr(torch_scatter, "segment_" + reduce + "_csr")
      jit = torch.jit.script(fn)
      out1 = fn(src, indptr)
      out2 = jit(src, indptr)
      if isinstance(out1, tuple):
          out1, arg_out1 = out1
          out2, arg_out2 = out2
          arg_expected = tensor(
              test["arg_" + reduce], torch.long, device
          )
          assert torch.all(arg_out1 == arg_expected)
          assert arg_out1.tolist() == arg_out2.tolist()
      assert torch.all(out1 == expected)
      assert out1.tolist() == out2.tolist()
  
      fn = getattr(torch_scatter, "segment_" + reduce + "_coo")
      jit = torch.jit.script(fn)
      out1 = fn(src, index)
      out2 = jit(src, index)
      if isinstance(out1, tuple):
          out1, arg_out1 = out1
          out2, arg_out2 = out2
          arg_expected = tensor(
              test["arg_" + reduce], torch.long, device
          )
          assert torch.all(arg_out1 == arg_expected)
          assert arg_out1.tolist() == arg_out2.tolist()
      assert torch.all(out1 == expected)
      assert out1.tolist() == out2.tolist()
  
file_path: pytorch_scatter/test/test_segment.py
imports: from itertools import product
  import pytest
  import torch
  import torch_scatter
  from torch.autograd import gradcheck
  from torch_scatter.testing import devices, dtypes, reductions, tensor
label_desc: Use image node_id 1 for calling a global function with example usage: test_forward(test, reduce, dtype, device) without return types