code: def _test_common(
      self,
      reduction,
      device,
      dtype,
      unsafe,
      axis,
      initial_value,
      data_arr,
      lengths_arr,
      expected_arr,
      expected_grad_arr,
      check_backward,
      lengths_dtype=torch.int,
  ):
      lengths = torch.tensor(
          lengths_arr, device=device, dtype=lengths_dtype
      )
      # generate offsets from lengths
      zeros_shape = list(lengths.shape)
      zeros_shape[-1] = 1
      offsets = torch.cat(
          (lengths.new_zeros(zeros_shape), lengths), -1
      ).cumsum_(-1)
  
      data = torch.tensor(
          data_arr,
          device=device,
          dtype=dtype,
          requires_grad=True,
      )
      expected_result = torch.tensor(
          expected_arr, device=device, dtype=dtype
      )
      expected_grad = torch.tensor(
          expected_grad_arr, device=device, dtype=dtype
      )
      for mode in ["lengths", "offsets"]:
          segment_reduce_kwargs = dict(
              axis=axis, unsafe=unsafe, initial=initial_value
          )
          if mode == "lengths":
              segment_reduce_kwargs["lengths"] = lengths
          else:
              segment_reduce_kwargs["offsets"] = offsets
          actual_result = torch._segment_reduce(
              data=data, reduce=reduction, **segment_reduce_kwargs
          )
          self.assertEqual(
              expected_result,
              actual_result,
              rtol=1e-02,
              atol=1e-05,
              equal_nan=True,
          )
  
          if not check_backward:
              return
  
          # Test backward
          actual_result.sum().backward()
          self.assertEqual(
              expected_grad,
              data.grad,
              rtol=1e-02,
              atol=1e-05,
              equal_nan=True,
          )
          data = data.clone().detach().requires_grad_(True)
  
          # gradcheck does not work well with bfloat16 or fp16 cpu types
          # also there is small numerical difference with fp32
          if dtype not in [torch.half, torch.bfloat16, torch.float]:
              # gradcheck does not like "nan" input, setting to random 10
              d_non_nan = np.nan_to_num(data_arr, nan=10)
              new_data = torch.tensor(
                  # [10 if v == float("nan") else v for v in data],
                  d_non_nan,
                  device=device,
                  dtype=dtype,
                  requires_grad=True,
              )
              self.assertTrue(
                  gradcheck(
                      lambda x: torch._segment_reduce(
                          data=x,
                          reduce=reduction,
                          **segment_reduce_kwargs,
                      ),
                      (new_data,),
                  )
              )
  
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 1 for calling the TestSegmentReductions obj's underlying member method code with example usage: obj._test_common(reduction, device, dtype, unsafe, axis, initial_value, data_arr, lengths_arr, expected_arr, expected_grad_arr, check_backward, lengths_dtype) without return types