Add Wishart distribution. (#1779)

* Add Wishart distribution.

* Reduce dimensionality for bijection tests of positive definite matrices.

* Add `WishartCholesky` distribution and use it as base for `Wishart`.

* Promote instead of broadcast Wishart parameters.

* Assert exactly one of parameters is specified and update shape inference.

* Implement `infer_shapes` for `Wishart` and `WishartCholesky`.

* Add entropy for Wishart distribution.

* Add sampled entropy test for distribution without `scipy` equivalent.

* Simplify `logabsdet` evaluation of `scale_tril`.

* Remove default `None` argument for concentration of Wishart distribution.

* Add `tri_logabsdet` utility function.

Author: tillahoffmann

docs/source/distributions.rst

@@ -380,6 +380,22 @@ Weibull
     :show-inheritance:
     :member-order: bysource
 
+Wishart
+^^^^^^^
+.. autoclass:: numpyro.distributions.continuous.Wishart
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+WishartCholesky
+^^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.continuous.WishartCholesky
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 ZeroSumNormal
 ^^^^^^^^^^^^^
 .. autoclass:: numpyro.distributions.continuous.ZeroSumNormal

numpyro/distributions/__init__.py

@@ -47,6 +47,8 @@
     StudentT,
     Uniform,
     Weibull,
+    Wishart,
+    WishartCholesky,
     ZeroSumNormal,
 )
 from numpyro.distributions.copula import GaussianCopula, GaussianCopulaBeta
@@ -194,6 +196,8 @@
     "Unit",
     "VonMises",
     "Weibull",
+    "Wishart",
+    "WishartCholesky",
     "ZeroInflatedDistribution",
     "ZeroInflatedPoisson",
     "ZeroInflatedNegativeBinomial2",