Implement the Stop-Gradient Differentiable Particle Filter

[DanWaxman]

This PR implements the stop-gradient differentiable PF (https://arxiv.org/abs/2106.10314) as a resampling decorator. It supersedes #202 after the merging of #207.

The following changes were made:

• Fix differentiability of the PF implementation
  • Previously, computing the gradient of the MLL would result in RunTime errors, as the default _inverse_cdf implementation on CPU used a jax.pure_callback. Since we don't expect the inverse_cdf function to be differentiable, anyways, we can safely wrap its inputs with a lax.stop_gradient call to ensure at least biased gradients are available.
• Introduce a Stop Gradient Resampling Decorator
• Adds a DPF example to the documentation
• Adds unit testing for the stop gradient PF

The resulting DPF is shown to provide strong score estimates in a toy example, now in the documentation.

To use the stop-gradient resampling, one simply uses the stop_gradient_decorator:

resampling_fn = (
    stop_gradient.stop_gradient_decorator(systematic.resampling)
    if differentiable_resampling
    else systematic.resampling
)
# In either case, do adaptive resampling
resampling_fn = adaptive.ess_decorator(resampling_fn, ess_threshold)

Unit testing in this setting is somewhat difficult, for a few reasons:

1. The estimator is still stochastic, and requires lots of trials or particles to make accurate. But this seems rather expensive for CI.
2. The existing utilities for generating test cases are time-varying, and provide somewhat ill-conditioned systems. For example, the the scale of observation noise as generated in cuthbertlib.kalman.generate can often be on the order of $10^-3$, which can make gradients pretty poorly behaved.

The resolution here is to choose an extremely simple system (random walk on x, identity Gaussian observation on y), and run with 10,000 particles. The score is estimated at the true parameter for the dynamics, and tested against the score from the ground-truth Kalman filter, taking the median over 10 runs, with a 20% relative error tolerance. This seems to run fast enough, at least locally. But happy to explore other options here.

[AdrienCorenflos]

I don't understand why it's hard to unittest this: it can be done almost without randomness.
Take a Gaussian weights w_n = N(0;x_n, \sigma^2) evaluated on a grid xs = np.linspace(-1, 1, 100). Compute the gradient of logsumexp(logws) before and after resampling and it should be the same if the wrapper is there and not otherwise? Am I missing something?

[DanWaxman]

Hi @AdrienCorenflos! Thanks for the suggestion! Indeed, that's a nice unit test -- I had "testing the stop-gradient DPF" in my head instead of "testing the stop gradient decorator," which is why I had written the test that way before. I've updated the unit tests to perform your version of the test (and reverted the test_particle_filters file).

[DanWaxman]

Responded to this round of feedback. Namely, I migrated the submodule name to cuthbertlib.resampling.autodiff and propagated those changes, and removed the erroneous printing statement. Comments regarding the end-to-end test in test_particle_filters.py were moot, since that test got replaced by the simpler unit test, but let me know if you think it's better to reinstate the end-to-end test as well.