def variable_synchronizer()

in lm_human_preferences/utils/core.py [0:0]

• 13 lines of code
• 4 McCabe index (conditional complexity)

def variable_synchronizer(comm, vars, *, limit=1<<28):
    """Synchronize `vars` from the root to other processs"""
    if comm.Get_size() == 1:
        return tf.no_op()

    # Split vars into chunks so that no chunk is over limit bytes
    batches = chunk_tensors(sorted(vars, key=lambda v: v.name), limit=limit)

    # Synchronize each batch, using a separate communicator to ensure safety
    prev = tf.no_op()
    for batch in batches:
        with tf.control_dependencies([prev]):
            assigns = []
            values = map_flat_bits(partial(mpi_bcast, comm), batch)
            for var, value in zip(batch, values):
                assigns.append(var.assign(value))
            prev = tf.group(*assigns)
    return prev