host_callback_depricated
When updating jax to 0.4.28, the experimental.host_callback becomes depricated. As far as I can see only terminal output on the minimization state is affected. I suggest replacing the calls to jax.exprimental.host_callback by jax.debug.callback.