search

patrick-kidger / equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Apache License 2.0
2.12k stars 142 forks source link

Hopefully compatibility with JAX 0.4.32 #856

Closed patrick-kidger closed 2 months ago

patrick-kidger commented 2 months ago

Possible fix for #854, let's see if tests pass.