search

rtqichen / torchdiffeq

Differentiable ODE solvers with full GPU support and O(1)-memory backpropagation.
MIT License
5.61k stars 930 forks source link

ComplexFloat implementation #213

Open itrosen opened 2 years ago

itrosen commented 2 years ago

Is there a way to use vectors and diff eqs with complex coefficients? Thanks!

a = torch.FloatTensor([[1,0],[0,1]])
y0 = torch.FloatTensor([1, 0])
def f(t, y):
        return -1.j * torch.matmul(a, y)
t_list = torch.linspace(0, 1, 11)
odeint(f, y0, t_list)

RuntimeError: "lt_cpu" not implemented for 'ComplexFloat'

pierreguilmin commented 1 year ago

I would love to see this feature as well!

orodrigoaraizabravo commented 1 year ago

This feature would be very much appreciated in the quantum computing community

rtqichen commented 1 year ago

I've just enabled complex types on the latest commit. Give it a try and if there are issues, let me know.

To install the latest commit:

pip install git+https://github.com/rtqichen/torchdiffeq
pierreguilmin commented 1 year ago

Thanks a lot!! 🥳 I tested it quickly on an example I had at hand, and it worked perfectly 👌 I'll make sure to notify you of any issues that I come across.

orodrigoaraizabravo commented 1 year ago

Thank you! This worked for me too! Just a heads up that I had to uninstall the package before installing the newest commit successfully. Use pip uninstall torchdiffeq