Thanks for the kind words! We've been doing a lot more work in this direction too, for both compiler-based automatic parallelization [0] and a work-in-progress

Thanks for the kind words! We've been doing a lot more work in this direction too, for both compiler-based automatic parallelization [0] and a work-in-progress pmap successor for 'manual' parallelism (per-device code and explicit collectives) [1] which composes seamlessly with the compiler-based stuff. [0] https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html https://jax.readthedocs.io/en/latest/notebooks/Distributed_a... [1] https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...