Change `jax.experimental.stax` to `jax.example_libraries.stax` (JAX 0.2.25) (!982) · Merge requests · Vicentini Filippo / netket

Vicentini Filippo requested to merge jax0.2.25 into master Nov 11, 2021

Created by: femtobit

This updates our references to stax, which has been moved to jax.example_libraries.stax in the latest JAX release (see #981).

Change `jax.experimental.stax` to `jax.example_libraries.stax` (JAX 0.2.25)