Incompatibility with jax 0.4.27

patrick-kidger / optimistix

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/

Apache License 2.0

310 stars 15 forks source link

Open tjltjl opened 4 months ago

tjltjl commented 4 months ago

Running the basic example produces:

when using the dependencies

Switching to 0.4.26 on both jax and jaxlib removes the error

patrick-kidger commented 4 months ago

Thanks for the report! Looks like an upstream JAX bug. I've opened https://github.com/google/jax/issues/21116.