Add some missing special functions to the documentation.
Add some missing special functions to the documentation.

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