[WIP] Add discrete/continuous support+more samplers

[rrkarim]

addr #187 #171
PR should include:

• Support for more samplers (the design should be discussed)
• Support for compound step sampling (maybe it is hard one to implement because of over-reliance on batching in tfp)
• Support for discrete var sampling
• Support auto_assign_method based on a set of certain criterias (like in pymc3)
• More reach trace for compound step
• More tests (they suck for now)

Current user design of compound step for user looks like this:

# the easiest one with automatic method assignment
# for each free var
model = ... # we define the model
pm.sample(model(), sampler_type="compound")
# user defines the methods for selected vars in model
from pymc4.mcmc.samplers import NUTS, RandomWalkM
pm.sample(model(), 
    sampler_type="compound", 
    sampler_methods=[
        ("var", RandomWalkM), 
        ("mean", NUTS)
    ], 
    xla=True)

There is currently support for categorical distribution mcmc sampling with custom state function, implemented in pymc4/distributions/state_functions.py, but more distributions should be tested and supported.

• Categorical
• **should be added as its tested and supported

[review-notebook-app]

Check out this pull request on 

You'll be able to see Jupyter notebook diff and discuss changes. Powered by ReviewNB.

[rrkarim]

There are some problems with sampling discrete distributions with xla enabled. I will look at the issue closer.

[rrkarim]

@junpenglao can you please provide some feedback on the issues #233, I've implemented the sampling of the models with categorical distribution (with support of new_state_fn) but the implementation still relies on CompounStep. I've just putted the tfp implementation of compound step in this branch. There is a brief notebook that will test sampling.

Also tests are not passing because of the notebooks, now they should be edited too, but it is going to be messy here.

[rrkarim]

Sorry for some notebooks edited, tests do not pass without it. There are also some issues with notebook testing, the radon notebook doesn't work on my local machine.

[rrkarim]

So basically, I had two ways of designing sampler methods. This PR implementation is similar to pymc3, which for me is more clear from user perspective. Another idea was to have MCMCSampler class with all the ops like sample, run_chains etc. and sampler subclasses that defined the parameters like kernel, adaptive_kernel, trace_fn and etc. Then _assign_default_methods method should be defined in MCMCSampler class (it is defined now in CompoundStep) so that user doesn't pass sampler_type argument, but the samplers are assigned automatically to each free var according to certain criterias. Basically, all the sampler methods are now the type of compound step. What I don't like about this design choice is how the parametrization for each sampler method is handled.

For now there is a lack of support for clever auto_assign_method which could assign appropriate methods based on set of criterias. And also, the trace for compound step is returning only target_log_prob, but ideally it could return mapping of bootstrap_results for each kernel used in compound step. I'm planning to implement them after some feedback will be given on current status, since I'm not sure if the implemented design is good enough.

[fonnesbeck]

I assume Unititled1.ipynb is not supposed to be there.