Document experimental.jet

Hi all,

From a lot of discussion in issues, it seems like jet is able to “compute higher order derivatives efficiently”. This functionality is particularly of my interest, but I can’t seem to find any more documentation than those discussion and a few tests. Could someone explain very shortly what is the arguments to jet(f, primals, series), in particular, what series represents?

Thanks!

[mattjj]

cc @duvenaud @jacobjinkelly @jessebett

I think @jessebett and @jacobjinkelly planned to write some kind of a post and/or docs for jet. In the meantime, there is the preprint Learning Differential Equations that are Easy to Solve, though that’s mostly about an application and doesn’t offer a ton in the way of documentation. I also wrote up some notes on the math and Sec 3 might help explain the reason the jet API looks how it does (but the most important bit for the internal implementation is Sec 4, and that’s a bit hand-wavey).

I think jet only becomes more efficient than jvp / vjp at third and higher orders though, so it might not be relevant to the Newton-Raphson application.

[pnkraemer]

Hi! I agree that it would be great to have some extra documentation for jet and would be willing to help. I started a corresponding discussion in #9536, because I wasn’t sure how good a place this issue here is to discuss details 😃