1. Formatting -- black style!
2. After profiling, implemented a better workaround for the
tfd.Categorical bug.  Replaced with an inherited Categorical2 class,
with overloaded log_prob function.
3. Fixed a  bug in EventTimeProposal where tfd.FiniteDiscrete was
being used on a non-strictly-increasing outcomes vector.
4. Increased number of meta-populations to 149 to mirror UK UTLA regions.

1. Fixed bug in covid.impl.event_time_mh._reverse_move where Multinomial one_hot tensor wasn't being advanced for the reverse move.  This is a temporary workaround while tf.nn.sparse_softmax_cross_entropy_with_logits is fixed.

2. Batched EventTimesProposal to move multiple meta-populations.

3. Batched _move_events to allow for multiple meta-population updates

4. Refactored covid.impl.event_time_proposal._abscumdiff so we calculate `bound_times` outside the function.  This allows clipping of times proposed outside [0, T] to [0, T], thus avoiding an error in tf.gather_nd in CPU mode.  The output of EventTimesProposal.sample and EventTimesProposal.log_prob might be inconsistent with what we expect, but that's okay as the move gets rejected as an out-of-bounds proposal anyway.

Avoid overflow bug in computing free_events

1. Added log_acceptance_probability to kernel results
2. Added test for XLA sparse_softmax issue.

1. Edited _move_events to take into account metapopulation selection.
2. Replace bound_t.ndim with bound_t.shape.rank for TF Graph Mode running.
3. Revised index calculation in `_abscumdiff`.  Re-wrote test to accommodate.
4. Code rationalisation in FilteredEventTimesProposal to return a simpler data structure.
5. Changed dimension of results.extra field to trace `x_star`.
6. Edited tests in `tests.test_mcmc.py` for new `_abscumdiff` interface.

1. Wrote wrapper `FilteredEventTimesProposal` for EventTimesProposal to subset metapopulations.  This currently selects just one meta-population, but the plan is to select more.
2. Moved test fixture to folder inside tests.
3. Wrote (failing) TestFilteredEventTimesProposal test.

Pair programming changes from 2020-06-12

- flatten out the distribution of delta_t; no longer a
JointDistributionNamed, just a FiniteDiscrete
- replaced OneHotCategorical with a Categorical (sample needs an expand dims
after)
- fix bug -- were using zeros instead of int max's for upper bounds
- thread dtypes through various functions
- inlined some arg-free distributions in JointDistributionNamed
- minor reformatting to appease Chris Suter's formatting compulsions
- pprint in tests; fix test data paths (weren't working for Chris
Suter's setup)

Further refactor of UncalibratedEventTimesUpdate.one_step, and solve of bug in reverse move calculation.

Added a "medium" simulation, which assumes 10 meta-populations of 1000 individuals each.

Implemented custom nametuple for tracing innerd of UncalibratedEventTimesUpdate

1. Corrected bug in UncalibratedEventTimesUpdate.one_step, and factored out some propagation functions.
2. Inserted a tf.cond into event times update step to deal with the case where there is no preceding or next event (e.g. S->E in an SEIR model doesn't have a preceding event).  This has reduced performance, and could possibly be replaced with a NOOP calculation instead.
3. Removed call to MH_within_Gibbs -- this proved too restrictive when requiring more than one update to a state_part per sweep of the MHwG sampler.

Use casting and tf.gather to avoid end-to-end XLA failure

Added inf checker to proposal, added mechanism to detect frequency of rejects due to -inf in likelihood.

1. Removed TransformedTransitionKernel.  This kernel transforms the log_prob function, so it it no longer compatible with the general target_log_prob.

2. Replaced with an UncalibratedLogRandomWalk class which is composed inside a MetropolisHastings kernel.

1. Implemented Haario-style covariance update for parameters
2. Fixed a bug in the use of underlying tfp.mcmc framework where TransformedTransitionKernel.bootstrap_results was returning incorrect results.

First pass implementation of data augmentation MCMC.  This is slow, buggy, and possibly not correct yet!  Use with caution!

Working event time update.  Not compatible with XLA due to tf.boolean_mask which depends on tf.where.