Comment by fhchl

Very cool! How does the optimized Pytoch code compare to the Jax implementation?