def METHOD_NAME(x, *, token=None, comm=MPI_jax_comm):
    """
    Computes the elementwise logical OR of an array or a scalar across all MPI
    processes, effectively equivalent to an elementwise any
    Args:
        a: The input array.
        token: An optional token to impose ordering of MPI operations
    Returns:
        out: The reduced array.
        token: an output token
    """
    if n_nodes == 1:
        return x, token
    else:
        import mpi4jax
        return mpi4jax.allreduce(x, op=MPI.MAX, comm=comm, token=token)