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

More fixes for JAX 0.4.34 #871

Closed patrick-kidger closed 1 month ago

patrick-kidger commented 1 month ago

Fixes https://github.com/patrick-kidger/diffrax/issues/508