code: def test_pytorch_scatter_test_cases(self, device, dtypes, reduce):
      val_dtype, length_dtype = dtypes
      # zero-length segments are filled with reduction inits contrary to pytorch_scatter.
      tests = [
          {
              "src": [1, 2, 3, 4, 5, 6],
              "index": [0, 0, 1, 1, 1, 3],
              "indptr": [0, 2, 5, 5, 6],
              "sum": [3, 12, 0, 6],
              "prod": [2, 60, 1, 6],
              "mean": [1.5, 4, float("nan"), 6],
              "min": [1, 3, float("inf"), 6],
              "max": [2, 5, -float("inf"), 6],
          },
          {
              "src": [
                  [1, 2],
                  [3, 4],
                  [5, 6],
                  [7, 8],
                  [9, 10],
                  [11, 12],
              ],
              "index": [0, 0, 1, 1, 1, 3],
              "indptr": [0, 2, 5, 5, 6],
              "sum": [[4, 6], [21, 24], [0, 0], [11, 12]],
              "prod": [[3, 8], [315, 480], [1, 1], [11, 12]],
              "mean": [
                  [2, 3],
                  [7, 8],
                  [float("nan"), float("nan")],
                  [11, 12],
              ],
              "min": [
                  [1, 2],
                  [5, 6],
                  [float("inf"), float("inf")],
                  [11, 12],
              ],
              "max": [
                  [3, 4],
                  [9, 10],
                  [-float("inf"), -float("inf")],
                  [11, 12],
              ],
          },
          {
              "src": [[1, 3, 5, 7, 9, 11], [2, 4, 6, 8, 10, 12]],
              "index": [[0, 0, 1, 1, 1, 3], [0, 0, 0, 1, 1, 2]],
              "indptr": [[0, 2, 5, 5, 6], [0, 3, 5, 6, 6]],
              "sum": [[4, 21, 0, 11], [12, 18, 12, 0]],
              "prod": [[3, 315, 1, 11], [48, 80, 12, 1]],
              "mean": [
                  [2, 7, float("nan"), 11],
                  [4, 9, 12, float("nan")],
              ],
              "min": [
                  [1, 5, float("inf"), 11],
                  [2, 8, 12, float("inf")],
              ],
              "max": [
                  [3, 9, -float("inf"), 11],
                  [6, 10, 12, -float("inf")],
              ],
          },
          {
              "src": [
                  [[1, 2], [3, 4], [5, 6]],
                  [[7, 9], [10, 11], [12, 13]],
              ],
              "index": [[0, 0, 1], [0, 2, 2]],
              "indptr": [[0, 2, 3, 3], [0, 1, 1, 3]],
              "sum": [
                  [[4, 6], [5, 6], [0, 0]],
                  [[7, 9], [0, 0], [22, 24]],
              ],
              "prod": [
                  [[3, 8], [5, 6], [1, 1]],
                  [[7, 9], [1, 1], [120, 143]],
              ],
              "mean": [
                  [[2, 3], [5, 6], [float("nan"), float("nan")]],
                  [[7, 9], [float("nan"), float("nan")], [11, 12]],
              ],
              "min": [
                  [[1, 2], [5, 6], [float("inf"), float("inf")]],
                  [[7, 9], [float("inf"), float("inf")], [10, 11]],
              ],
              "max": [
                  [[3, 4], [5, 6], [-float("inf"), -float("inf")]],
                  [[7, 9], [-float("inf"), -float("inf")], [12, 13]],
              ],
          },
          {
              "src": [[1, 3], [2, 4]],
              "index": [[0, 0], [0, 0]],
              "indptr": [[0, 2], [0, 2]],
              "sum": [[4], [6]],
              "prod": [[3], [8]],
              "mean": [[2], [3]],
              "min": [[1], [2]],
              "max": [[3], [4]],
          },
          {
              "src": [[[1, 1], [3, 3]], [[2, 2], [4, 4]]],
              "index": [[0, 0], [0, 0]],
              "indptr": [[0, 2], [0, 2]],
              "sum": [[[4, 4]], [[6, 6]]],
              "prod": [[[3, 3]], [[8, 8]]],
              "mean": [[[2, 2]], [[3, 3]]],
              "min": [[[1, 1]], [[2, 2]]],
              "max": [[[3, 3]], [[4, 4]]],
          },
      ]
      for test in tests:
          data = torch.tensor(
              test["src"],
              dtype=val_dtype,
              device=device,
              requires_grad=True,
          )
          indptr = torch.tensor(
              test["indptr"], dtype=length_dtype, device=device
          )
          dim = indptr.ndim - 1
          # calculate lengths from indptr
          lengths = torch.diff(indptr, dim=dim)
          expected = torch.tensor(
              test[reduce], dtype=val_dtype, device=device
          )
  
          actual_result = torch._segment_reduce(
              data=data,
              reduce=reduce,
              lengths=lengths,
              axis=dim,
              unsafe=True,
          )
          self.assertEqual(actual_result, expected)
  
          # test offsets
          actual_result = torch._segment_reduce(
              data=data,
              reduce=reduce,
              offsets=indptr,
              axis=dim,
              unsafe=True,
          )
          self.assertEqual(actual_result, expected)
  
          if val_dtype == torch.float64:
  
              def fn(x, mode="lengths"):
                  initial = 1
                  # supply initial values to prevent gradcheck from failing for 0 length segments
                  # where nan/inf are reduction identities that produce nans when calculating the numerical jacobian
                  if reduce == "min":
                      initial = 1000
                  elif reduce == "max":
                      initial = -1000
                  segment_reduce_args = {x, reduce}
                  segment_reduce_kwargs = dict(
                      axis=dim, unsafe=True, initial=initial
                  )
                  if mode == "lengths":
                      segment_reduce_kwargs[mode] = lengths
                  elif mode == "offsets":
                      segment_reduce_kwargs[mode] = indptr
                  return torch._segment_reduce(
                      *segment_reduce_args, **segment_reduce_kwargs
                  )
  
              self.assertTrue(
                  gradcheck(
                      partial(fn, mode="lengths"),
                      (data.clone().detach().requires_grad_(True)),
                  )
              )
              self.assertTrue(
                  gradcheck(
                      partial(fn, mode="offsets"),
                      (data.clone().detach().requires_grad_(True)),
                  )
              )
  
file_path: pytorch/test/test_segment_reductions.py
imports: from itertools import product
  from functools import partial
  import numpy
  import torch
  from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes
  from torch.testing._internal.common_utils import TestCase, run_tests, gradcheck, parametrize, skipIfRocm
label_desc: Use image node_id 5 for calling the TestSegmentReductions obj's underlying member method code with example usage: obj.test_pytorch_scatter_test_cases(device, dtypes, reduce) and returns: torch