Import side effect: GPU preallocatoin

[Jacob-Stevens-Haas]

At import, diffrax appears to grab the full preallocation of GPU memory even when it won't ever be used on device. This is hard for me as a downstream library developer because I can't assume anything about backend in the library. When it comes to testing, I have in the past relied on a pytest fixture to determine how to set the default device. But since these are after imports, they can't be used with diffrax

Setting the device in jax config at import time works, but it feels like bad practice and wont allow generic control of longer-running benchmarks which may be run on CPU or GPU.

This may be due to the same issue in lineax - of your packages, only diffrax and lineax grab the GPU as an import side effect.

diffrax version: 0.7.0
python version 3.12.3

[lockwo]

Does it grab 100% or ~75%? Since in general JAX does a lot of preallocation (see https://docs.jax.dev/en/latest/gpu_memory_allocation.html), so this might a JAX thing (that can be fixed with those XLA flags)

[Jacob-Stevens-Haas]

It is the standard jax preallocation, but AFAIK XLA flags run into the same issues as jax.config management, in that they need to be set before the first allocation. That preallocation happens for diffrax (maybe because of lineax?) at import time, which isn't ideal. Library code shouldn't set environment variables for the user, and I'd like to be able to write tests that defer choice of compute device until later.

In general, if I'm running tests while debugging or running an interactive session, I want tests on the CPU so I don't cause OOM errors by too many processes. OTOH, I also want to be able to run some of the same tests on GPU, and choosing devices at import time creates another import-time side effect.

[patrick-kidger]

I'm surprised that Diffrax causes this at import time. We're generally careful to avoid performing e.g. jnp.array calls at the global scope, which is the usual thing that causes the JAX backend to be initialized.

If we've missed anything (maybe bisect the Diffrax codebase with a test script of just import diffrax) then I'd be happy to take a PR on this.

[Jacob-Stevens-Haas]

I'll give it a swing!