When using pmap I often miss a parallel mean function lax.pmean(x, axis_name).
I guess it's not too complicated to write lax.psum / axis_size. However, there is currently also no way to retrieve the size of a parallel axis within a parallel computation.