Clarify docstrings regarding usage of static arguments in jit and vmap. (#3484)
Clarify docstrings regarding usage of static arguments in jit and vmap. (#3484)

The docstring for pmap does not currently mention that any "non-data"
arguments need to be indicated in `static_broadcasted_argnums`. This
commit updates the docs to parallel those for `jax.jit` which does
explain this. Additionally, a remark is added to the `static_*` argument
descriptions on both jit and pmap, so that this point can be understood
without reading the whole docstring.
google/jax,tensorflow/probability,google/jax,google/jax,tensorflow/probability,google/jax