Implements relaxed inference for compartmental models

[fritzo]

Addresses #2426
Replaces #2510

This implements an asymptotically exact relaxed inference strategy for CompartmentalModel whereby:

1. An auxiliary variable is introduced such that the flow constraints are always satisfied (however constraints due to observations may still be violated, so SVI still doesn't work and SMC initialization is still required)
2. Rather than enumerating, gradients are computed using a differentiably_quantize() function that quantizes but keeps gradients. This avoids the cost of enumeration which is exponential in number of compartments, and also avoids coupling through time so that sparse observations are cheaper (do not require an additional auxiliary variable).

This method differs from bouncy ball MCMC (Nimshimura, Dunson, Lu 2017) in that we still use an L2 Gaussian kinetic energy rather than an L1 Laplace kinetic energy, thus permitting reuse of existing implementations of HMC and NUTS. And despite the gradients in (2) being approximate, this method is still asymptotically correct because the exact density is used (and hence the MH correction compensates for bad gradients). Indeed for Binomial potentials the gradients are quite accurate (see notebook):

Mixing Issues

In practice this mixes poorly. I believe the cause is inability to propagate gradients to avoid feasibility boundaries. Experiments show that early in learning trajectories become infeasible, and HMC step size is reduced to very small.

$ python examples/contrib/epidemiology/sir.py -p 1000 -d 10 -f 10 -R0 3 --relax --plot
Simulating from a SimpleSIRModel
Observed 23/43 infections in population of 1000:
1 0 1 0 3 2 1 3 5 7
INFO 	 Running inference...
Warmup:   0%|                                                                              | 0/300 [00:00, ?it/s]INFO 	 Heuristic init: R0=1.9, rho=0.949
DEBUG infeasible
DEBUG infeasible
DEBUG infeasible
DEBUG infeasible
DEBUG infeasible
DEBUG infeasible
Warmup:   0%|                                     | 1/300 [00:00,  3.90it/s, step size=7.47e-02, acc. prob=0.013]DEBUG infeasible
DEBUG infeasible
Warmup:   6%|██                                  | 17/300 [00:03,  4.34it/s, step size=7.11e-04, acc. 

Tested

• smoke tests
• example tests
• verify mixing on larger examples

[fritzo]

Closing until we can resolve mixing issues.