Tests failing when running with JAX v0.6.0 due to breaking changes #299
Description
Various tests failed in the last scheduled run with an error message
AttributeError: register_custom_call_target has been removed in JAX v0.6.0; use the JAX FFI instead
This is due to JAX v0.6.0 being released on April 16th 2025 which introduced some breaking changes. I believe #290 will address the issue in the longer term by using the JAX foreign function interface, but for now we may want to do a quickfix of pinning to jax<0.6.0