update for jax 0.2.14, mpi4jax 0.3.1 (!753) · Merge requests · Vicentini Filippo / netket

Vicentini Filippo requested to merge pv/jax-version into master Jun 11, 2021

Created by: PhilipVinc

I'll soon tag a new mpi4jax release (0.3.1) that works with the latest jax as well as the old ones. Since there are a few features that require more recent versions of mpi4jax (like vs.to_array()) I bumped the minimum required mpi4jax version to the latest.

This PR also changes the code used to check the lowest compatible mpi4jax version to use a custom function instead of depending on setuptools, which will simplify the conda installation.

update for jax 0.2.14, mpi4jax 0.3.1

Merge request reports