Importing diffrax initializes Jax backend

[brosand]

Importing diffrax for the first time initializes a jax backend per this jax file. This raises a warning if there is no GPU on the machine, which can cause issues with a CI pipeline that uses Diffrax, as hiding this warning would prevent the warning from showing up later. Would there be an easy fix to stop this backend initialization when importing Diffrax? I don't know much about this level of the code.

[patrick-kidger]

So I think this should be fixed on main, but not yet pushed to PyPI. Can you install directly from GitHub and see if this resolves your issue? If it does then I'd be happy to do a new release ASAP.

[brosand]

Yes this does solve it, thanks! A new release would be great, so we can merge into our own repo without messing up CI. Glad to resolve this so fast.

[patrick-kidger]

New version should appear on PyPI in the next few hours, once the auto-deploy/test scripts have done their job.