Support PyTorch JIT compilation

[fritzo]

What would it take to support JIT-compilation of an entire SVI.step()?

For motivation, consider a beta bernoulli model which is 10x slower in Pyro than Edward.

Feasibility

PyTorch @jit.compile appears to be functional input->output. Therefore we would need to at least

1. run a warm-up iteration to populate the param store
2. expose the Pyro param store to the jit compiler
3. expose a helper .differentiable_loss() inside loss_and_grads (everything up to .backward())

Tasks

• assess feasibility
• Support torch.jit in ELBO implementations #1109 refactor ELBO to expose .differentiable_loss() (see also Migrate from .backward() to torch.autograd.grad() #628)
• Add helper decorator for JITing parametrized functions #1126 add an abstract jit decorator to make it easy to write new jitted functions
• Update JIT usage to PyTorch 0.4.1 #1276 Change to using torch.jit.trace instead of the deprecated torch.jit.compile
• [JIT] Support torch.distributions.utils.broadcast_all() pytorch/pytorch#10041 Support broadcast_all() in the jit
• [distributions] Make more distributions jittable pytorch/pytorch#10321 Support Categorical.log_prob() in the jit
• Add xfailing jit tests to many anchor examples #1296 add --jit option to SVI examples
• Use JIT traced potential energy computation in HMC #1299 support jit in HMC
• add tutorial

Remaining PyTorch issues

To see remaining PyTorch jit issues:

1. Install PyTorch master

   cd ~/github/pytorch/pytorch  # or wherever your clone lives
   git checkout master
   git pull
   python setup.py build develop

2. Run Pyro's jit tests with the pytorch-0.4.1 branch of Pyro

   cd ~/github/uber/pyro  # or wherever your clone lives
   git checkout pytorch-0.4.1
   git pull
   make install
   make test-jit  # saves results to jit.log

Here is an example jit.log.

[eb8680]

Are any torch.distributions samplers or log-probabilities supported by the JIT at all? Does the JIT have the same broadcasting semantics?

[fritzo]

I believe this is closed by #1431