Fix callbacks

[stefanocortinovis]

Checklist

• [ x ] I've formatted the new code by running hatch run dev:format before committing.
• [ NA ] I've added tests for new code.
• [ NA ] I've added docstrings for the new code.

Description

Migrate from jax.experimental.host_callback() to jax.debug.callback() (see here). Fixes #479.

Issue Number: #479

[github-actions]

Thank you for opening your first PR into GPJax!

If you have not heard from us in a while, please feel free to ping @gpjax/developers or anyone who has commented on the PR. Most of our reviewers are volunteers and sometimes things fall through the cracks.

You can also join us on Slack for real-time discussion.

For details on testing, writing docs, and our review process, please see the developer guide

We strive to be a welcoming and open project. Please follow our Code of Conduct.

[daniel-dodd]

Thanks @stefanocortinovis!

[thomaspinder]

Thanks for fixing this so swiftly @stefanocortinovis. Pending successful tests, could you please bump this line to 0.9.3?

[stefanocortinovis]

Hi @daniel-dodd, @thomaspinder! I've opened a couple of other PRs. Happy to bump the version once you are happy with the changes!