astralord/jax_parallel

Supplementary code for https://astralord.github.io/posts/exploring-parallel-strategies-with-jax/

Supplementary code for Exploring Parallel Strategies with Jax