fixes for jax 0.3.3 (!1141) · Merge requests · Vicentini Filippo / netket

Vicentini Filippo requested to merge pv/jax3 into master Mar 18, 2022

Created by: PhilipVinc

The latest jax released earlier today/yesterday broke some small things because it removed jax.ops.index.

fixes for jax 0.3.3