Add CirculantNormal distribution.

[tillahoffmann]

This PR adds a CirculantNormal distribution which is a multivariate normal distribution where the covariance matrix has circular boundary conditions. It uses the Fourier transform introduced in #1762 to evaluate the log probability, scaling with n * log(n) instead of the usual n ** 3. The PR

• Implements the CirculantNormal distribution.
• Adds a PackRealFastFourierCoefficientsTransform transform. This is required to transform a real vector of size n to a complex vector of size n // 2 + 1 which can be passed to the rfft method. See Add complex constraint and real Fourier transform. #1762 (comment) for a brief initial discussion.
• Adds a positive_definite_circulant_vector constraint. This is really just required for the test/test_distributions.py::test_distribution_constraints tests to pass.
• Adds an example notebook, illustrating how the CirculantNormal can be used to accelerate Gaussian process inference.
• Fix the Jacobian of the Fourier transform (which turns out not to be the identity but depends on the norm parameter).

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[juanitorduz]

It would be nice to put the legends outside the plots as in https://stackoverflow.com/questions/4700614/how-to-put-the-legend-outside-the-plot :)

Reply via ReviewNB

[tillahoffmann]

I'm broadcasting to the full shape here instead of just promoting the shapes. That's to guarantee that covariance_row always has shape batch_shape + event_shape. If we only promoted shapes, we might end up with a shape (1, 1, n) if the input covariance_row has shape (n,) but the input loc has shape (a, b, n).

As an aside, this may also be relevant to MultivariateNormal distribution where covariance_matrix may not have the right batch dimensions.

>>> from jax import numpy as jnp
>>> from numpyro.distributions import MultivariateNormal
>>> 
>>> d = MultivariateNormal(jnp.zeros((3, 4, 5)), jnp.eye(5))
>>> d
<numpyro.distributions.continuous.MultivariateNormal object at 0x145cc7d00 with 
batch shape (3, 4) and event shape (5,)>
>>> d.covariance_matrix.shape
(1, 1, 5, 5)  # Expected (3, 4, 5, 5)

[tillahoffmann]

I think the failing tests are due to issues with the data download and not related to the changes in this PR. There are a few further failing tests due to numerics (maybe associated with a new jax release?).