Thanks for your comment - I wanted to keep it as general as possible, and I think speed touches basically every possible use of JAX, but I definitely hear what

Thanks for your comment - I wanted to keep it as general as possible, and I think speed touches basically every possible use of JAX, but I definitely hear what you are saying. For DL, jvps and quick/easy Hessians (along with speed) are really the reasons to use JAX imo. I like your comment on thinking about implementing models/thinking about them as parameterized solutions