[On dev branch] Tuple rewrite

lxuechen commented 3 years ago

Contents in this PR:

Remove support for tuple y0 and tuple-based sde for sdeint.
- Rationale: This feature is somewhat niche. Supporting this would likely require adding an alternative tuple-based ForwardSDE that has the correct versions of f, g_prod, gdg_prod, and dg_ga_jvp (per our email discussion).
Make AdjointSDE take in ForwardSDE and use the latter's g_prod.
- This simplifies g_prod of the former.
Make AdjointSDE take in (and output) flattened tensors and perform flattening/unflattening inside the registered functions.
- One advantage of this is that we only need to unflatten 2 tensors now.
Refactor solvers to not be based on tuples.
Explicit flattening/unflattening in adjoint.py
- This differs slightly from the current implementation in torchdiffeq, where the _TupleFunc construction is in _check_inputs.
- Rationale: We assume that sdeint.integrate only supports tensor-based y0 and sde, and we should enforce this usage for now.
Filter out the parameters that don't require gradients in sdeint_adjoint.
Fix plot folder names for files in diagnostics.
Temporarily deprecating logqp support.
- This should either be added back or the example latent_sde rewritten with state augmentation when we get the general noise adjoint working with the log-ODE scheme. (I view getting the log-ODE scheme-based adjoint working for general noise Strat SDE as a milestone. )
Fix scalar noise shape check and it's functionality in sdeint inconsistency. Hence, resolving #27.

The rate diagnostics files run, but I haven't been able to obtain rates due to #36. I need to run these when this is fixed. All tests pass.

Overall, this PR should also make adding dg_ga_jvp for the adjoint easier.

patrick-kidger commented 3 years ago

Also, tests/__init__.pyc shouldn't be included in the commit.

patrick-kidger commented 3 years ago

Overall, things are looking so much cleaner now. These changes make me very happy.

patrick-kidger commented 3 years ago

Just spotted, both sdeint and its adjoint have the docstring for names still specifying the possibility of prior drift.

lxuechen commented 3 years ago

Just spotted, both sdeint and its adjoint have the docstring for names still specifying the possibility of prior drift.

Done.

patrick-kidger commented 3 years ago

This PR looks good to me. I've not hit the merge button as I assume the plan is to merge #40 first, so that the diagnostics can be run.

lxuechen commented 3 years ago

Thanks for the reviews @patrick-kidger! All rates preserved after running the diagnostics. Merging now.

google-research / torchsde