Or Golan

Hi! What is missing for this to be merged? Is there any way I can help with it?

Hi, is this still an issue? I'm trying to use the GPU with pymanopt and I see all sorts of errors.

I'm having the exact same issue on Arch, with python 3.10 and installing through pip instead of conda.

From what I looked, it seems that this line: https://github.com/patrick-kidger/diffrax/blob/cec091c5e4cc4311f64ae3aa09a371db5fe766ee/diffrax/integrate.py#L771 ```python ys = jax.tree_map(lambda y: jnp.full((out_size,) + jnp.shape(y), jnp.inf), y0) ``` is what causing complex numbers to be cast as...

I was not able to continue with this PR (see the comment above).

It seems that `jax.experimental.host_callback.id_print` is doing nothing when I run in debug mode. Should I use it directly in a print statement?

I have the same issue as @ngm - after disabling `file-templates` and `snippets` it started working again. The issue started after I defined a few new snippets.

I have the exact issue with MacOS

I'd be happy to work on this with @Aaron-Robertson if he wants, or to work on it alone. I want to make sure I understand the scope of the work...

I haven't