search

google-research / torchsde

Differentiable SDE solvers with GPU support and efficient sensitivity analysis.
Apache License 2.0
1.52k stars 195 forks source link

Operate on tensors rather than tuple-of-tensors? #8

Closed patrick-kidger closed 3 years ago

patrick-kidger commented 3 years ago

Hi Xuechen - I appreciate I'm bombarding this repository with a lot of issues/PRs.

Something we've recently found over in torchdiffeq#115 is that it's noticeably more efficient to have the internals of the library operate on tensors, rather than tuples-of-tensors, and have the adjoint method simply torch.cat its different channels together.

Of course this also makes the code a whole lot cleaner, and I expect easier to develop going forward.

Not something I've got time to put together a PR for I'm afraid, but I thought I'd flag this up as something you might be interested in over here.

lxuechen commented 3 years ago

Hi Patrick,

Thanks for the note! It's interesting that using torch.cat for the output can speed things up by much. I definitely learned something here.

I'll try to patch this up over a weekend in the future, and feel free to open up other issues even if you don't have cycles to work on them. Really appreciate it!

patrick-kidger commented 3 years ago

No problem. More issues to follow!

patrick-kidger commented 3 years ago

Done in #37.