Also fix the docstring.

Previously one got the following errors:

mpi4jax:
```
MPI_ABORT was invoked on rank 0 in communicator MPI COMMUNICATOR 4 
CREATE FROM 0 with errorcode 10.
```
mpi4py:
``
mpi4py.MPI.Exception: MPI_ERR_OP: invalid reduce operation
```
reproducer:

```python
import netket as nk
import jax.numpy as jnp
x = 1.j
y = jnp.ones(2, dtype=complex)
#nk.utils.mpi.mpi_max(x)
#nk.utils.mpi.mpi_max_jax(x)
nk.utils.mpi.mpi_max_jax(y)
```