requiremented environment

[azh2000xy]

Hello, I am reproducing your code failed due to jax version issues and configured environment. Can you release the environment and corresponding version of the pips in this paper，such as numpy==....,jax==.....

[MilesCranmer]

Hey @azh2000xy,
It should work for the most recent version, with the change that jax.experimental -> jax.example_libraries

- from jax.experimental import stax
+ from jax.example_libraries import stax
- from jax.experimental import optimizers
+ from jax.example_libraries import optimizers

cc @greydanus - we should probably apply this update to the examples?

Cheers,
Miles

[MilesCranmer]

I added a note about this to the README in 2fc26cb.

[MilesCranmer]

I made a PR here: #4 to fix these. Want to take a look?

[healer-21]

Hello, can you disclose the data set you used？