How does nested pmap with axis_name work?

[llan-ml]

from functools import partial
import jax

@partial(pmap, axis_name='rows')
@partial(pmap, axis_name='cols')
def normalize(x):
  row_normed = x / jax.lax.psum(x, 'rows')
  col_normed = x / jax.lax.psum(x, 'cols')
  doubly_normed = x / jax.lax.psum(x, ('rows', 'cols'))
  return row_normed, col_normed, doubly_normed

In the above snippet, the two nesting pmaps refer to the row and column axes, respectively. How does this work? I did not find some useful documentation to explain the behaviors of nesting multiple pmaps. Can anyone help to give some more explanation? Thanks!