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_a...
[1] https://jax.readthedocs.io/en/latest/jep/14273-shard-map.htm...