add pmap docstring
add pmap docstring

Co-authored-by: Peter Hawkins <34b19541630131950b5f204f2c5f2c647a3a7475@google.com>

google/jax,tensorflow/probability,tensorflow/probability,google/jax,google/jax,google/jax