Add RecursiveLinearTransform for linear state space models.

[tillahoffmann]

This PR adds a RecursiveLinearTransform which is a linear transformation applied recursively such that $y_t = A y_{t - 1} + x_t$ for $t > 0$, where $x_t$ and $y_t$ are $p$-vectors and $A$ is a $p\times p$ transition matrix. The series is initialized by $y_0 = 0$.

This transform can be used to easily declare linear state space models, e.g., a Cauchy random walk is

>>> from jax import random
>>> from jax import numpy as jnp
>>> import numpyro
>>> from numpyro import distributions as dist
>>>
>>> def cauchy_random_walk():
...     return numpyro.sample(
...         "x",
...         dist.TransformedDistribution(
...             dist.Cauchy(0, 1).expand([10, 1]).to_event(1),
...             dist.transforms.RecursiveLinearTransform(jnp.eye(1)),
...         ),
...     )

A Kalman-style model for a rocket with state y = (position, velocity) is

>>> def rocket_trajectory():
...     scale = numpyro.sample(
...         "scale",
...         dist.HalfCauchy(1).expand([2]).to_event(1),
...     )
...     transition_matrix = jnp.array([[1, 1], [0, 1]])
...     return numpyro.sample(
...         "x",
...         dist.TransformedDistribution(
...             dist.Normal(0, scale).expand([10, 2]).to_event(1),
...             dist.transforms.RecursiveLinearTransform(transition_matrix),
...         ),
...     )

This PR also makes a few minor changes (happy to factor out if you prefer):

• Reformat the ExplicitReparam from Add explicit reparametrizer. #1754 to comply with the stricter linting.
• Add the RealFastFourierTransform from Add complex constraint and real Fourier transform. #1762 to the documentation.
• Ignore venv directory in the update_headers.py script.
• Add autogenerated documentation sources docs/source/{examples,tutorials,getting_started.rst} to .gitignore.
• Verify log_abs_det_jacobian implementation using autodiff.

[fehiepsi]

Maybe 'matrix' -> squared matrix for clarity. Currently, the bias x is time-dependent, but the matrix is constant. Do you plan to make the class name more verbose to reflect that?

[tillahoffmann]

Maybe 'matrix' -> squared matrix for clarity.

👍

Currently, the bias x is time-dependent, but the matrix is constant. Do you plan to make the class name more verbose to reflect that?

Sorry, I explained poorly. The x here is the argument of the transform rather than the bias in the AffineTransform, for example. Keeping the transition matrix constant means that the transform can be applied to sequences of arbitrary length.

[fehiepsi]

x is the (time-dependent) bias (or noise if we place a normal distribution over it) of a linear dynamical model. Typically, people use other notations, something like x_t = Ax{t-1} + b_t.

[tillahoffmann]

Fair. I was using the x/y notation here to stick with the typical arguments of the Transform classes. Do you have an idea for a more explanatory name?

[fehiepsi]

The name looks 👌 to me. :)

[fehiepsi]

I think you might want jnp.swapaxes(self.transition_matrix, -1, -2)

[tillahoffmann]

I've updated it to use einsum because we have shapes (..., p, p) for the transition matrices A and (..., n, p) for the states x. Inside the scan function, we are dealing with state of shape (..., p) because we're scanning along the n dimension. There may very well be a better way to do this.

[fehiepsi]

similarly, jnp.swapaxes(self.transition_matrix, -1, -2)

[fehiepsi]

sounds reasonable to me, this is sort of a shear transformation, so the Jacobian determinant is 1.

[tillahoffmann]

Yes, because of the temporal nature of the transform, the Jacobian is triangular. Because the x only appears additively, the diagonal is one, leading to the unit Jacobian.

[fehiepsi]

Thanks, @tillahoffmann!