def _rotate_tensor_block(buf): ntasks = mpi.pool.size tasks = list(range(ntasks)) tasks = tasks[rank:] + tasks[:rank] for task in tasks: if task != rank: buf = mpi.rotate(buf) yield task, buf
def rotate(): buf_prefetch[0] = mpi.rotate(buf, blocking=False)