# Return a decorator that runs the function only if rank == 0,
# otherwise waits for a broadcast from rank 0
def broadcast_from_zero(rank, comm, root=0):
    def decorator(func):
        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            if comm is None:
                return func(*args, **kwargs)

            if rank == root:
                val = func(*args, **kwargs)
            else:
                val = None
            comm.bcast(val, root=root)
            return val
        return wrapped
    return decorator