Add a pyro.contrib.cevae module (#2202)

* Refactor and add docs to pyro.contrib.cevae

* Add smoke tests

* Add link to docs

* Fix typos

* Install graphviz on travis

* Add a jit_trace method

* Improve jit support; move example

* Rename method

* Get CEVAE to work on large real dataset

* Support batching in CEVAE.ite()

* Generalize to Bernoulli or Normal outcomes

* Fix clamping bounds

* Add StudentT observation distribution

* Add Exponential outcome distribution

* Add LaplaceNet

* Fix architecture according to original code

* Address review comment

* Fix placement of activations in twin networks

* Fix sign errors

.travis.yml

@@ -16,6 +16,7 @@ addons:
             - gcc-8
             - g++-8
             - ninja-build
+            - graphviz  # for sphinx docs
 
 cache:
     directories:

docs/source/conf.py

@@ -42,6 +42,7 @@
     'sphinx.ext.ifconfig',  #
     'sphinx.ext.viewcode',  #
     'sphinx.ext.githubpages',  #
+    'sphinx.ext.graphviz',  #
     'sphinx.ext.autodoc',
     'sphinx.ext.doctest',
 ]

docs/source/contrib.cevae.rst

@@ -0,0 +1,82 @@
+Causal Effect VAE
+=================
+
+.. automodule:: pyro.contrib.cevae
+
+CEVAE Class
+-----------
+.. autoclass:: pyro.contrib.cevae.CEVAE
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+CEVAE Components
+----------------
+.. autoclass:: pyro.contrib.cevae.Model
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.Guide
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.TraceCausalEffect_ELBO
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+Utilities
+---------
+.. autoclass:: pyro.contrib.cevae.FullyConnected
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.DistributionNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.BernoulliNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.ExponentialNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.LaplaceNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.NormalNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.StudentTNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+.. autoclass:: pyro.contrib.cevae.DiagNormalNet
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource

docs/source/index.rst

@@ -31,6 +31,7 @@ Pyro Documentation
 
    contrib.autoname
    contrib.bnn
+   contrib.cevae
    contrib.easyguide
    contrib.examples
    contrib.gp