-
Notifications
You must be signed in to change notification settings - Fork 624
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
GPU-enabled solver? #2320
Comments
It's in development. qutip-jax is not up to date with the latest versions. We expect an official release of these features in March. |
Thank you! I'll try qutip-jax. Looking forward to the official release! |
Is there a simplest case I can run the |
I suggest you to use the example in qutip/qutip-jax#26. First ensure that jax is using the gpu:
After I would suggest to use the branch in qutip/qutip-jax#20.
|
Just to add, I tried to benchmark Eric's jax data layer a bit more with an Ising model, the example is at the end of this colab notebook we made for a tutorial talk, which shows some crossover in performance at certain system sizes: https://colab.research.google.com/drive/1RcgX7oEzGjzPAF8Ryus54Q5UmyMddmLA?usp=sharing Note colab does not have free GPUs, so you will have to download and use it locally. also in the actual ising example, replace |
Hi again, I followed the instruction in this doc qutip-jax, and run the following codes:
and this will output a warning:
Is it a thing if I need a relatively good simulation? |
It should not affect the results of simulations. The diffrax package does not interact with complex number directly when used through qutip-jax. Gradient or other derivatives could be affected, this need more testing. |
Thanks. I come across another problem when I try to run the following codes on a L4 GPU provided by google colab:
And an error will occur:
So I assume the qutip-jax doesn't work with L4 GPU, right? What GPU or TPU can be used in this case? |
I don't know, since jax is developed by google I expect it to work well with gpus it provide through collab... We don't use schur decomposition for mesolve directly. I think it's the integrator from diffrax that does, if not it could be the norm (it's using trace norm which call sqrtm instead of trace, fixed in #2408). Maybe trying other ODE solver or not normalizing would work. Neill seems to have run most of his test using dopri:
|
Thank you!, it does work. But there is another problem regarding solving time-dependent Hamiltonian evolution in qutip-jax:
And this will give:
Any other way we can solve the evolution with time-dependent Hamiltonian? |
Only function decorated with jit are supported as coefficient in with diffrax. |
Describe the Issue!
Hi, I would like to know if there is any way I can run the solver (mesove, mcsolve, etc.) in a GPU-enabled setting? I don't see any tutorial or manual to support this feature.
The text was updated successfully, but these errors were encountered: