def mpi_op()

in spinup/utils/mpi_tools.py [0:0]

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

def mpi_op(x, op):
    x, scalar = ([x], True) if np.isscalar(x) else (x, False)
    x = np.asarray(x, dtype=np.float32)
    buff = np.zeros_like(x, dtype=np.float32)
    allreduce(x, buff, op=op)
    return buff[0] if scalar else buff