jax

585cefc8 - document pmap with pytrees, fixes #1486