Having read some Jax high-performance code, I do like Jax, but it does feel a bit too abstract and low level sometimes. Maybe there aren’t good coding conventions or performance trumped them? Definitely needs improvement on error messages, as well.

For example, a long chain of pmaps, each with some sort of device partitioning logic, not JIT compiling is extremely hard to understand. I basically had to binary search code until the compile errors disappeared.