Supplementary code for Exploring Parallel Strategies with Jax
nvbkdw / jax_parallel Goto Github PK
View Code? Open in Web Editor NEWThis project forked from astralord/jax_parallel
Supplementary code for https://astralord.github.io/posts/exploring-parallel-strategies-with-jax/