4 ms·
And 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 ov
by PartiallyTyped 5y ago
And 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