3 ms·
Yes exactly! This is huge. Hessian optimization is really easy with JAX, haven't tried it in Julia though
by SleekEagle 5y ago
Yes exactly! This is huge. Hessian optimization is really easy with JAX, haven't tried it in Julia though
- PartiallyTyped 5y agoAnd very fast given that you compile the procedure! I am considering writing an article on this and posting it here because I have seen enormous improvements over non jitted code, and that excluded jax.vmap.
- SleekEagle 5y agoThere's a comparison of JAX with PyTorch for Hessian calculation here! https://www.assemblyai.com/blog/why-you-should-or-shouldnt-be-using-jax-in-2022/#hessians https://www.assemblyai.com/blog/why-you-should-or-shouldnt-b... Would definitely be interested in an article like that if you decide to write it
- ChrisRackauckas 5y agoHere's Hessian-Free Newton-Krylov on neural ODEs with Julia: https://diffeqflux.sciml.ai/dev/examples/second_order_adjoints/ https://diffeqflux.sciml.ai/dev/examples/second_order_adjoin... . It's just standard tutorial stuff at this point.