Add some doc strings to lax primitives.
Add some doc strings to lax primitives.

Since lax is a semipublic API, its public methods need at least minimal documentation. Many of the docstrings added in this PR are somewhat redundant, but at least a few contain useful information, and the documentation reads better with at least some minimal text for each function.

Hide some methods that shouldn't be public from the lax API docs.

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