def worker(rank, data, expect, port): if mge.get_device_count("gpu") < world_size: return dist.init_process_group("localhost", port, world_size, rank, rank) inp = tensor(data) output = all_reduce_min(inp) assert np.allclose(output.numpy(), expect)
def worker(data, expect): rank = dist.get_rank() inp = tensor(data[rank]) output = all_reduce_min(inp) assert np.allclose(output.numpy(), expect[rank])