#1 fix
Merge branch 'stable-devel' of https://github.com/jerrymf/JAX2 into stable-devel