code: def get_partition_weights(weights, parts):
      """Return the amount of weight in each partition."""
      costs = [0] * (len(parts) - 1)
      P = len(parts) - 1
      for p in range(P):
          start = parts[p]
          stop = parts[p + 1]
          costs[p] = sum(weights[start:stop])
      return costs
  
imports: import pytest
  import torch
  import deepspeed.comm
  from deepspeed.runtime.utils import partition_uniform
  from deepspeed.runtime.utils import partition_balanced
  from deepspeed.runtime.utils import prefix_sum_inc
  from deepspeed.runtime.utils import PartitionedTensor
  from unit.common import DistributedTest
label_desc: Use image node_id 2 for calling a global function with example usage: get_partition_weights(weights, parts) and returns: costs