def mpi_allreduce_sum()

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

• 15 lines of code
• 3 McCabe index (conditional complexity)

def mpi_allreduce_sum(values, *, comm):
    if comm.Get_size() == 1:
        return values
    orig_dtype = values.dtype
    if hvd is None:
        orig_shape = values.shape
        def _allreduce(vals):
            buf = np.zeros(vals.shape, np.float32)
            comm.Allreduce(vals, buf, op=MPI.SUM)
            return buf
        values = tf.py_func(_allreduce, [values], tf.float32)
        values.set_shape(orig_shape)
    else:
        values = hvd.mpi_ops._allreduce(values)
    return tf.cast(values, dtype=orig_dtype)