fix netket tools info broken by mpi4jax

Created by: PhilipVinc

Before I was unsafely relying on some mpi4jax internals. Now i use the exported function if available otherwise I check before accessing internals...

Should be merged only after https://github.com/mpi4jax/mpi4jax/pull/109 is merged, in case the name is changed.