code: def unsorted_segment_sum(
      data: torch.Tensor, segment_ids: torch.Tensor, num_segments: int
  ) -> torch.Tensor:
      """Computes the sum along segments of a tensor. Analogous to tf.unsorted_segment_sum.
  
      Parameters
      ----------
      data: torch.Tensor
          A tensor whose segments are to be summed.
      segment_ids: torch.Tensor
          The segment indices tensor.
      num_segments: int
          The number of segments.
  
      Returns
      -------
      tensor: torch.Tensor
  
      Examples
      --------
      >>> segment_ids = torch.Tensor([0, 1, 0]).to(torch.int64)
      >>> data = torch.Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [4, 3, 2, 1]])
      >>> num_segments = 2
      >>> result = unsorted_segment_sum(data=data,
      ...                               segment_ids=segment_ids,
      ...                               num_segments=num_segments)
      >>> data.shape[0]
      3
      >>> segment_ids.shape[0]
      3
      >>> len(segment_ids.shape)
      1
      >>> result
      tensor([[5., 5., 5., 5.],
              [5., 6., 7., 8.]])
  
      """
  
      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."
          )
  
      s = torch.prod(torch.tensor(data.shape[1:])).long()
      segment_ids = segment_ids.repeat_interleave(s).view(
          segment_ids.shape[0], *data.shape[1:]
      )
  
      # data.shape and segment_ids.shape should be equal
      assert data.shape == segment_ids.shape
      shape: List[int] = [num_segments] + list(data.shape[1:])
      tensor: torch.Tensor = torch.zeros(*shape).scatter_add(
          0, segment_ids, data.float()
      )
      tensor = tensor.type(data.dtype)
      return tensor
  
file_path: deepchem/deepchem/utils/pytorch_utils.py
imports: import scipy
  import torch
  from typing import Callable, Union, List, Generator, Tuple
label_desc: Use image node_id 2 for calling a global function with example usage: unsorted_segment_sum(data, segment_ids, num_segments) and returns: tensor