[Feature]: Compatibility issues with jax and flax

[Saisai-Xia]

Feature Request Type

Documentation

Have you searched existing issues?

Yes

Is your feature request related to a problem?

Hi, I tried to reproduce the results from the paper, but there are some library compatibility issues in the environment configuration file provided in this repo, especially with JAX and Flax. Could you please provide the specific compatible environment configuration?

Describe features you want to add to SPU

A clear and concise description of what you want to happen.

Describe features you want to add to SPU

A clear and concise description of any alternative solutions or features you've considered.

[fionser]

@Saisai-Xia Please try again.