Return type is tuple of solution and nfe not just solution

[FedericoV]

• https://github.com/jacobjinkelly/easy-neural-ode/blob/1ae0836959fe61d11b09540369f3dcd230e58565/lib/ode.py#L554-558

The return type is actually a tuple of the solution and the number of function evaluations:

easy-neural-ode/lib/ode.py

Line 747 in 1ae0836

return jax.vmap(unravel)(out), nfe

Super easy to clean up, I just thought it was some issues with pytrees the first time I came across it.

I'm currently going through the repo to understand the code base a bit better: would you prefer if I opened an issue for every tiny thing I find, or would it be easier to work with a general clean up PR?

[jacobjinkelly]

Hi @FedericoV, good catch, thanks for bringing this up! There are probably a number of spots where the docstring is different from what the code actually does, my apologies.

If you're able to, one PR for any general clean-up you find would be great! But of course feel free to open several PRs/issues as you see fit if the topics are distinct. If anything is confusing or unclear, definitely err on the side of making an issue for it instead of spending a lot of time trying to figure it out, as it's probably a silly mistake on my end :p

[FedericoV]

No apology necessary at all, and thank you for making this code all open. I was trying to understand the function signature of jax.jet a little bit better, and I was struggling a bit since most of the example in the code are fairly complex and nested pretty deep. Do you have any minimal example of using jet to calculate the higher order gradients of a simple function?

…

On Wed, Jan 20, 2021 at 5:26 PM Jacob Kelly ***@***.***> wrote: Hi @FedericoV <https://github.com/FedericoV>, good catch, thanks for bringing this up! There are probably a number of spots where the docstring is different from what the code actually does, my apologies. If you're able to, one PR for any general clean-up you find would be great! But of course feel free to open several PRs/issues as you see fit if the topics are distinct. If anything is confusing or unclear, definitely err on the side of making an issue for it instead of spending a lot of time trying to figure it out, as it's probably a silly mistake on my end :p — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAEZ25TM6EYNBRYQ43H44RLS257C3ANCNFSM4WLEPWKQ> .

[jacobjinkelly]

I would recommend having a look at the tests for jet which can be found here