code: def test_segment_sum():
      data = torch.Tensor([[1, 2, 3, 4], [4, 3, 2, 1], [5, 6, 7, 8]])
      segment_ids = torch.Tensor([0, 0, 1]).to(torch.int64)
  
      if len(segment_ids.shape) != 1:
          raise AssertionError("segment_ids have be a 1-D tensor")
  
      if data.shape[0] != segment_ids.shape[0]:
          raise AssertionError(
              "segment_ids should be the same size as dimension 0 of input."
          )
  
      result = dc.utils.pytorch_utils.segment_sum(
          data=data, segment_ids=segment_ids
      )
  
      assert np.allclose(
          np.array(result),
          np.load("deepchem/utils/test/assets/result_segment_sum.npy"),
          atol=1e-04,
      )
  
file_path: deepchem/deepchem/utils/test/test_pytorch_utils.py
imports: import deepchem
  import numpy
  import pytest
label_desc: Use image node_id 2 for calling a global function with example usage: test_segment_sum() without return types