Function get_partition_start

mpu/random.py:236–240

(item)

def get_partition_start(item):
    global mp_rank, mp_size, mp_group
    partition_size = get_partition_size(item)
    start = partition_size * mp_rank
    return int(start)

def get_partition_size(item):
    global mp_rank, mp_size, mp_group