Comment by jdeaton

The abstractions provided by JAX for parallelism are beautiful. JAX is an absolute master-class in programming-interface design and a lesson in the power of providing composable primitive operations and FP inspired design. An astounding amount of complexity is hidden from the user behind primitives like pmap, and the power is exposed in such a simple interface.