We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
DMM Example  Model and guide dims disagree when IAF is applied
When the number of IAF is non zero, the latent dimension changes to the same value as the batch size.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pyro/distributions/transformed_distribution.py`
Content:
```
1 from __future__ import absolute_import, division, print_function
2 
3 import numbers
4 from abc import ABCMeta, abstractmethod
5 
6 import torch
7 import torch.nn as nn
8 from six import add_metaclass
9 from torch.autograd import Variable
10 
11 from pyro.distributions.distribution import Distribution
12 from pyro.distributions.util import copy_docs_from
13 from pyro.nn import AutoRegressiveNN
14 
15 
16 @copy_docs_from(Distribution)
17 class TransformedDistribution(Distribution):
18     """
19     Transforms the base distribution by applying a sequence of `Bijector`s to it.
20     This results in a scorable distribution (i.e. it has a `log_pdf()` method).
21 
22     :param base_distribution: a (continuous) base distribution; samples from this distribution
23         are passed through the sequence of `Bijector`s to yield a sample from the
24         `TransformedDistribution`
25     :type base_distribution: pyro.distribution.Distribution
26     :param bijectors: either a single Bijector or a sequence of Bijectors wrapped in a nn.ModuleList
27     :returns: the transformed distribution
28     """
29 
30     def __init__(self, base_distribution, bijectors, *args, **kwargs):
31         super(TransformedDistribution, self).__init__(*args, **kwargs)
32         self.reparameterized = base_distribution.reparameterized
33         self.base_dist = base_distribution
34         if isinstance(bijectors, Bijector):
35             self.bijectors = nn.ModuleList([bijectors])
36         elif isinstance(bijectors, nn.ModuleList):
37             for bijector in bijectors:
38                 assert isinstance(bijector, Bijector), \
39                     "bijectors must be a Bijector or a nn.ModuleList of Bijectors"
40             self.bijectors = bijectors
41 
42     def sample(self, *args, **kwargs):
43         """
44         :returns: a sample y
45         :rtype: torch.autograd.Variable
46 
47         Sample from base distribution and pass through bijector(s)
48         """
49         x = self.base_dist.sample(*args, **kwargs)
50         next_input = x
51         for bijector in self.bijectors:
52             y = bijector(next_input)
53             if bijector.add_inverse_to_cache:
54                 bijector._add_intermediate_to_cache(next_input, y, 'x')
55             next_input = y
56         return next_input
57 
58     def batch_shape(self, x=None, *args, **kwargs):
59         return self.base_dist.batch_shape(*args, **kwargs)
60 
61     def event_shape(self, *args, **kwargs):
62         return self.base_dist.event_shape(*args, **kwargs)
63 
64     def log_prob(self, y, *args, **kwargs):
65         """
66         :param y: a value sampled from the transformed distribution
67         :type y: torch.autograd.Variable
68 
69         :returns: the score (the log pdf) of y
70         :rtype: torch.autograd.Variable
71 
72         Scores the sample by inverting the bijector(s) and computing the score using the score
73         of the base distribution and the log det jacobian
74         """
75         value = y
76         log_det_jacobian = 0.0
77         for bijector in reversed(self.bijectors):
78             log_det_jacobian += bijector.batch_log_det_jacobian(value, *args, **kwargs)
79             value = bijector.inverse(value)
80         base_log_pdf = self.base_dist.log_prob(value, *args, **kwargs)
81         if not isinstance(log_det_jacobian, numbers.Number):
82             log_det_jacobian = log_det_jacobian.contiguous().view(*base_log_pdf.size())
83             assert log_det_jacobian.size() == base_log_pdf.size(), \
84                 'Invalid batch_log_det_jacobian().size():\nexpected {}\nactual {}'.format(
85                         base_log_pdf.size(), log_det_jacobian.size())
86         return base_log_pdf - log_det_jacobian
87 
88 
89 @add_metaclass(ABCMeta)
90 class Bijector(nn.Module):
91     """
92     Abstract class `Bijector`. `Bijector` are bijective transformations with computable
93     log det jacobians. They are meant for use in `TransformedDistribution`.
94     """
95 
96     def __init__(self, *args, **kwargs):
97         super(Bijector, self).__init__(*args, **kwargs)
98         self.add_inverse_to_cache = False
99 
100     @abstractmethod
101     def __call__(self, *args, **kwargs):
102         """
103         Virtual forward method
104 
105         Invokes the bijection x=>y
106         """
107         raise NotImplementedError()
108 
109     @abstractmethod
110     def inverse(self, *args, **kwargs):
111         """
112         Virtual inverse method
113 
114         Inverts the bijection y => x.
115         """
116         raise NotImplementedError()
117 
118     def log_det_jacobian(self, *args, **kwargs):
119         """
120         Default logdet jacobian method.
121 
122         Computes the log det jacobian `|dy/dx|`
123         """
124         return self.batch_log_det_jacobian(*args, **kwargs).sum()
125 
126     def batch_log_det_jacobian(self, *args, **kwargs):
127         """
128         Virtual elementwise logdet jacobian method.
129 
130         Computes the log abs det jacobian `|dy/dx|`
131         """
132         raise NotImplementedError()
133 
134 
135 class InverseAutoregressiveFlow(Bijector):
136     """
137     An implementation of an Inverse Autoregressive Flow. Together with the `TransformedDistribution` this
138     provides a way to create richer variational approximations.
139 
140     Example usage::
141 
142     >>> base_dist = Normal(...)
143     >>> iaf = InverseAutoregressiveFlow(...)
144     >>> pyro.module("my_iaf", iaf)
145     >>> iaf_dist = TransformedDistribution(base_dist, iaf)
146 
147     Note that this implementation is only meant to be used in settings where the inverse of the Bijector
148     is never explicitly computed (rather the result is cached from the forward call). In the context of
149     variational inference, this means that the InverseAutoregressiveFlow should only be used in the guide,
150     i.e. in the variational distribution. In other contexts the inverse could in principle be computed but
151     this would be a (potentially) costly computation that scales with the dimension of the input (and in
152     any case support for this is not included in this implementation).
153 
154     :param input_dim: dimension of input
155     :type input_dim: int
156     :param hidden_dim: hidden dimension (number of hidden units)
157     :type hidden_dim: int
158     :param sigmoid_bias: bias on the hidden units fed into the sigmoid; default=`2.0`
159     :type sigmoid_bias: float
160     :param permutation: whether the order of the inputs should be permuted (by default the conditional
161         dependence structure of the autoregression follows the sequential order)
162     :type permutation: bool
163 
164     References:
165 
166     1. Improving Variational Inference with Inverse Autoregressive Flow [arXiv:1606.04934]
167     Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling
168 
169     2. Variational Inference with Normalizing Flows [arXiv:1505.05770]
170     Danilo Jimenez Rezende, Shakir Mohamed
171 
172     3. MADE: Masked Autoencoder for Distribution Estimation [arXiv:1502.03509]
173     Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle
174     """
175 
176     def __init__(self, input_dim, hidden_dim, sigmoid_bias=2.0, permutation=None):
177         super(InverseAutoregressiveFlow, self).__init__()
178         self.input_dim = input_dim
179         self.hidden_dim = hidden_dim
180         self.arn = AutoRegressiveNN(input_dim, hidden_dim, output_dim_multiplier=2, permutation=permutation)
181         self.sigmoid = nn.Sigmoid()
182         self.sigmoid_bias = Variable(torch.Tensor([sigmoid_bias]))
183         self._intermediates_cache = {}
184         self.add_inverse_to_cache = True
185 
186     def get_arn(self):
187         """
188         :rtype: pyro.nn.AutoRegressiveNN
189 
190         Return the AutoRegressiveNN associated with the InverseAutoregressiveFlow
191         """
192         return self.arn
193 
194     def __call__(self, x, *args, **kwargs):
195         """
196         :param x: the input into the bijection
197         :type x: torch.autograd.Variable
198 
199         Invokes the bijection x=>y; in the prototypical context of a TransformedDistribution `x` is a
200         sample from the base distribution (or the output of a previous flow)
201         """
202         hidden = self.arn(x)
203         sigma = self.sigmoid(hidden[:, 0:self.input_dim] + self.sigmoid_bias.type_as(hidden))
204         mean = hidden[:, self.input_dim:]
205         y = sigma * x + (Variable(torch.ones(sigma.size())).type_as(sigma) - sigma) * mean
206         self._add_intermediate_to_cache(sigma, y, 'sigma')
207         return y
208 
209     def inverse(self, y, *args, **kwargs):
210         """
211         :param y: the output of the bijection
212         :type y: torch.autograd.Variable
213 
214         Inverts y => x. As noted above, this implementation is incapable of inverting arbitrary values
215         `y`; rather it assumes `y` is the result of a previously computed application of the bijector
216         to some `x` (which was cached on the forward call)
217         """
218         if (y, 'x') in self._intermediates_cache:
219             x = self._intermediates_cache.pop((y, 'x'))
220             return x
221         else:
222             raise KeyError("Bijector InverseAutoregressiveFlow expected to find" +
223                            "key in intermediates cache but didn't")
224 
225     def _add_intermediate_to_cache(self, intermediate, y, name):
226         """
227         Internal function used to cache intermediate results computed during the forward call
228         """
229         assert((y, name) not in self._intermediates_cache),\
230             "key collision in _add_intermediate_to_cache"
231         self._intermediates_cache[(y, name)] = intermediate
232 
233     def batch_log_det_jacobian(self, y, *args, **kwargs):
234         """
235         Calculates the elementwise determinant of the log jacobian
236         """
237         if (y, 'sigma') in self._intermediates_cache:
238             sigma = self._intermediates_cache.pop((y, 'sigma'))
239         else:
240             raise KeyError("Bijector InverseAutoregressiveFlow expected to find" +
241                            "key in intermediates cache but didn't")
242         log_sigma = torch.log(sigma)
243         if 'log_pdf_mask' in kwargs:
244             log_sigma = log_sigma * kwargs['log_pdf_mask']
245         return log_sigma.sum(-1)
246
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pyro/distributions/transformed_distribution.py b/pyro/distributions/transformed_distribution.py
--- a/pyro/distributions/transformed_distribution.py
+++ b/pyro/distributions/transformed_distribution.py
@@ -55,7 +55,10 @@
             next_input = y
         return next_input
 
-    def batch_shape(self, x=None, *args, **kwargs):
+    def shape(self, *args, **kwargs):
+        return self.base_dist.shape(*args, **kwargs)
+
+    def batch_shape(self, *args, **kwargs):
         return self.base_dist.batch_shape(*args, **kwargs)
 
     def event_shape(self, *args, **kwargs):
@@ -242,4 +245,4 @@
         log_sigma = torch.log(sigma)
         if 'log_pdf_mask' in kwargs:
             log_sigma = log_sigma * kwargs['log_pdf_mask']
-        return log_sigma.sum(-1)
+        return log_sigma

{"golden_diff": "diff --git a/pyro/distributions/transformed_distribution.py b/pyro/distributions/transformed_distribution.py\n--- a/pyro/distributions/transformed_distribution.py\n+++ b/pyro/distributions/transformed_distribution.py\n@@ -55,7 +55,10 @@\n             next_input = y\n         return next_input\n \n-    def batch_shape(self, x=None, *args, **kwargs):\n+    def shape(self, *args, **kwargs):\n+        return self.base_dist.shape(*args, **kwargs)\n+\n+    def batch_shape(self, *args, **kwargs):\n         return self.base_dist.batch_shape(*args, **kwargs)\n \n     def event_shape(self, *args, **kwargs):\n@@ -242,4 +245,4 @@\n         log_sigma = torch.log(sigma)\n         if 'log_pdf_mask' in kwargs:\n             log_sigma = log_sigma * kwargs['log_pdf_mask']\n-        return log_sigma.sum(-1)\n+        return log_sigma\n", "issue": "DMM Example  Model and guide dims disagree when IAF is applied\nWhen the number of IAF is non zero, the latent dimension changes to the same value as the batch size.\n", "before_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport numbers\nfrom abc import ABCMeta, abstractmethod\n\nimport torch\nimport torch.nn as nn\nfrom six import add_metaclass\nfrom torch.autograd import Variable\n\nfrom pyro.distributions.distribution import Distribution\nfrom pyro.distributions.util import copy_docs_from\nfrom pyro.nn import AutoRegressiveNN\n\n\n@copy_docs_from(Distribution)\nclass TransformedDistribution(Distribution):\n    \"\"\"\n    Transforms the base distribution by applying a sequence of `Bijector`s to it.\n    This results in a scorable distribution (i.e. it has a `log_pdf()` method).\n\n    :param base_distribution: a (continuous) base distribution; samples from this distribution\n        are passed through the sequence of `Bijector`s to yield a sample from the\n        `TransformedDistribution`\n    :type base_distribution: pyro.distribution.Distribution\n    :param bijectors: either a single Bijector or a sequence of Bijectors wrapped in a nn.ModuleList\n    :returns: the transformed distribution\n    \"\"\"\n\n    def __init__(self, base_distribution, bijectors, *args, **kwargs):\n        super(TransformedDistribution, self).__init__(*args, **kwargs)\n        self.reparameterized = base_distribution.reparameterized\n        self.base_dist = base_distribution\n        if isinstance(bijectors, Bijector):\n            self.bijectors = nn.ModuleList([bijectors])\n        elif isinstance(bijectors, nn.ModuleList):\n            for bijector in bijectors:\n                assert isinstance(bijector, Bijector), \\\n                    \"bijectors must be a Bijector or a nn.ModuleList of Bijectors\"\n            self.bijectors = bijectors\n\n    def sample(self, *args, **kwargs):\n        \"\"\"\n        :returns: a sample y\n        :rtype: torch.autograd.Variable\n\n        Sample from base distribution and pass through bijector(s)\n        \"\"\"\n        x = self.base_dist.sample(*args, **kwargs)\n        next_input = x\n        for bijector in self.bijectors:\n            y = bijector(next_input)\n            if bijector.add_inverse_to_cache:\n                bijector._add_intermediate_to_cache(next_input, y, 'x')\n            next_input = y\n        return next_input\n\n    def batch_shape(self, x=None, *args, **kwargs):\n        return self.base_dist.batch_shape(*args, **kwargs)\n\n    def event_shape(self, *args, **kwargs):\n        return self.base_dist.event_shape(*args, **kwargs)\n\n    def log_prob(self, y, *args, **kwargs):\n        \"\"\"\n        :param y: a value sampled from the transformed distribution\n        :type y: torch.autograd.Variable\n\n        :returns: the score (the log pdf) of y\n        :rtype: torch.autograd.Variable\n\n        Scores the sample by inverting the bijector(s) and computing the score using the score\n        of the base distribution and the log det jacobian\n        \"\"\"\n        value = y\n        log_det_jacobian = 0.0\n        for bijector in reversed(self.bijectors):\n            log_det_jacobian += bijector.batch_log_det_jacobian(value, *args, **kwargs)\n            value = bijector.inverse(value)\n        base_log_pdf = self.base_dist.log_prob(value, *args, **kwargs)\n        if not isinstance(log_det_jacobian, numbers.Number):\n            log_det_jacobian = log_det_jacobian.contiguous().view(*base_log_pdf.size())\n            assert log_det_jacobian.size() == base_log_pdf.size(), \\\n                'Invalid batch_log_det_jacobian().size():\\nexpected {}\\nactual {}'.format(\n                        base_log_pdf.size(), log_det_jacobian.size())\n        return base_log_pdf - log_det_jacobian\n\n\n@add_metaclass(ABCMeta)\nclass Bijector(nn.Module):\n    \"\"\"\n    Abstract class `Bijector`. `Bijector` are bijective transformations with computable\n    log det jacobians. They are meant for use in `TransformedDistribution`.\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super(Bijector, self).__init__(*args, **kwargs)\n        self.add_inverse_to_cache = False\n\n    @abstractmethod\n    def __call__(self, *args, **kwargs):\n        \"\"\"\n        Virtual forward method\n\n        Invokes the bijection x=>y\n        \"\"\"\n        raise NotImplementedError()\n\n    @abstractmethod\n    def inverse(self, *args, **kwargs):\n        \"\"\"\n        Virtual inverse method\n\n        Inverts the bijection y => x.\n        \"\"\"\n        raise NotImplementedError()\n\n    def log_det_jacobian(self, *args, **kwargs):\n        \"\"\"\n        Default logdet jacobian method.\n\n        Computes the log det jacobian `|dy/dx|`\n        \"\"\"\n        return self.batch_log_det_jacobian(*args, **kwargs).sum()\n\n    def batch_log_det_jacobian(self, *args, **kwargs):\n        \"\"\"\n        Virtual elementwise logdet jacobian method.\n\n        Computes the log abs det jacobian `|dy/dx|`\n        \"\"\"\n        raise NotImplementedError()\n\n\nclass InverseAutoregressiveFlow(Bijector):\n    \"\"\"\n    An implementation of an Inverse Autoregressive Flow. Together with the `TransformedDistribution` this\n    provides a way to create richer variational approximations.\n\n    Example usage::\n\n    >>> base_dist = Normal(...)\n    >>> iaf = InverseAutoregressiveFlow(...)\n    >>> pyro.module(\"my_iaf\", iaf)\n    >>> iaf_dist = TransformedDistribution(base_dist, iaf)\n\n    Note that this implementation is only meant to be used in settings where the inverse of the Bijector\n    is never explicitly computed (rather the result is cached from the forward call). In the context of\n    variational inference, this means that the InverseAutoregressiveFlow should only be used in the guide,\n    i.e. in the variational distribution. In other contexts the inverse could in principle be computed but\n    this would be a (potentially) costly computation that scales with the dimension of the input (and in\n    any case support for this is not included in this implementation).\n\n    :param input_dim: dimension of input\n    :type input_dim: int\n    :param hidden_dim: hidden dimension (number of hidden units)\n    :type hidden_dim: int\n    :param sigmoid_bias: bias on the hidden units fed into the sigmoid; default=`2.0`\n    :type sigmoid_bias: float\n    :param permutation: whether the order of the inputs should be permuted (by default the conditional\n        dependence structure of the autoregression follows the sequential order)\n    :type permutation: bool\n\n    References:\n\n    1. Improving Variational Inference with Inverse Autoregressive Flow [arXiv:1606.04934]\n    Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling\n\n    2. Variational Inference with Normalizing Flows [arXiv:1505.05770]\n    Danilo Jimenez Rezende, Shakir Mohamed\n\n    3. MADE: Masked Autoencoder for Distribution Estimation [arXiv:1502.03509]\n    Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle\n    \"\"\"\n\n    def __init__(self, input_dim, hidden_dim, sigmoid_bias=2.0, permutation=None):\n        super(InverseAutoregressiveFlow, self).__init__()\n        self.input_dim = input_dim\n        self.hidden_dim = hidden_dim\n        self.arn = AutoRegressiveNN(input_dim, hidden_dim, output_dim_multiplier=2, permutation=permutation)\n        self.sigmoid = nn.Sigmoid()\n        self.sigmoid_bias = Variable(torch.Tensor([sigmoid_bias]))\n        self._intermediates_cache = {}\n        self.add_inverse_to_cache = True\n\n    def get_arn(self):\n        \"\"\"\n        :rtype: pyro.nn.AutoRegressiveNN\n\n        Return the AutoRegressiveNN associated with the InverseAutoregressiveFlow\n        \"\"\"\n        return self.arn\n\n    def __call__(self, x, *args, **kwargs):\n        \"\"\"\n        :param x: the input into the bijection\n        :type x: torch.autograd.Variable\n\n        Invokes the bijection x=>y; in the prototypical context of a TransformedDistribution `x` is a\n        sample from the base distribution (or the output of a previous flow)\n        \"\"\"\n        hidden = self.arn(x)\n        sigma = self.sigmoid(hidden[:, 0:self.input_dim] + self.sigmoid_bias.type_as(hidden))\n        mean = hidden[:, self.input_dim:]\n        y = sigma * x + (Variable(torch.ones(sigma.size())).type_as(sigma) - sigma) * mean\n        self._add_intermediate_to_cache(sigma, y, 'sigma')\n        return y\n\n    def inverse(self, y, *args, **kwargs):\n        \"\"\"\n        :param y: the output of the bijection\n        :type y: torch.autograd.Variable\n\n        Inverts y => x. As noted above, this implementation is incapable of inverting arbitrary values\n        `y`; rather it assumes `y` is the result of a previously computed application of the bijector\n        to some `x` (which was cached on the forward call)\n        \"\"\"\n        if (y, 'x') in self._intermediates_cache:\n            x = self._intermediates_cache.pop((y, 'x'))\n            return x\n        else:\n            raise KeyError(\"Bijector InverseAutoregressiveFlow expected to find\" +\n                           \"key in intermediates cache but didn't\")\n\n    def _add_intermediate_to_cache(self, intermediate, y, name):\n        \"\"\"\n        Internal function used to cache intermediate results computed during the forward call\n        \"\"\"\n        assert((y, name) not in self._intermediates_cache),\\\n            \"key collision in _add_intermediate_to_cache\"\n        self._intermediates_cache[(y, name)] = intermediate\n\n    def batch_log_det_jacobian(self, y, *args, **kwargs):\n        \"\"\"\n        Calculates the elementwise determinant of the log jacobian\n        \"\"\"\n        if (y, 'sigma') in self._intermediates_cache:\n            sigma = self._intermediates_cache.pop((y, 'sigma'))\n        else:\n            raise KeyError(\"Bijector InverseAutoregressiveFlow expected to find\" +\n                           \"key in intermediates cache but didn't\")\n        log_sigma = torch.log(sigma)\n        if 'log_pdf_mask' in kwargs:\n            log_sigma = log_sigma * kwargs['log_pdf_mask']\n        return log_sigma.sum(-1)\n", "path": "pyro/distributions/transformed_distribution.py"}], "after_files": [{"content": "from __future__ import absolute_import, division, print_function\n\nimport numbers\nfrom abc import ABCMeta, abstractmethod\n\nimport torch\nimport torch.nn as nn\nfrom six import add_metaclass\nfrom torch.autograd import Variable\n\nfrom pyro.distributions.distribution import Distribution\nfrom pyro.distributions.util import copy_docs_from\nfrom pyro.nn import AutoRegressiveNN\n\n\n@copy_docs_from(Distribution)\nclass TransformedDistribution(Distribution):\n    \"\"\"\n    Transforms the base distribution by applying a sequence of `Bijector`s to it.\n    This results in a scorable distribution (i.e. it has a `log_pdf()` method).\n\n    :param base_distribution: a (continuous) base distribution; samples from this distribution\n        are passed through the sequence of `Bijector`s to yield a sample from the\n        `TransformedDistribution`\n    :type base_distribution: pyro.distribution.Distribution\n    :param bijectors: either a single Bijector or a sequence of Bijectors wrapped in a nn.ModuleList\n    :returns: the transformed distribution\n    \"\"\"\n\n    def __init__(self, base_distribution, bijectors, *args, **kwargs):\n        super(TransformedDistribution, self).__init__(*args, **kwargs)\n        self.reparameterized = base_distribution.reparameterized\n        self.base_dist = base_distribution\n        if isinstance(bijectors, Bijector):\n            self.bijectors = nn.ModuleList([bijectors])\n        elif isinstance(bijectors, nn.ModuleList):\n            for bijector in bijectors:\n                assert isinstance(bijector, Bijector), \\\n                    \"bijectors must be a Bijector or a nn.ModuleList of Bijectors\"\n            self.bijectors = bijectors\n\n    def sample(self, *args, **kwargs):\n        \"\"\"\n        :returns: a sample y\n        :rtype: torch.autograd.Variable\n\n        Sample from base distribution and pass through bijector(s)\n        \"\"\"\n        x = self.base_dist.sample(*args, **kwargs)\n        next_input = x\n        for bijector in self.bijectors:\n            y = bijector(next_input)\n            if bijector.add_inverse_to_cache:\n                bijector._add_intermediate_to_cache(next_input, y, 'x')\n            next_input = y\n        return next_input\n\n    def shape(self, *args, **kwargs):\n        return self.base_dist.shape(*args, **kwargs)\n\n    def batch_shape(self, *args, **kwargs):\n        return self.base_dist.batch_shape(*args, **kwargs)\n\n    def event_shape(self, *args, **kwargs):\n        return self.base_dist.event_shape(*args, **kwargs)\n\n    def log_prob(self, y, *args, **kwargs):\n        \"\"\"\n        :param y: a value sampled from the transformed distribution\n        :type y: torch.autograd.Variable\n\n        :returns: the score (the log pdf) of y\n        :rtype: torch.autograd.Variable\n\n        Scores the sample by inverting the bijector(s) and computing the score using the score\n        of the base distribution and the log det jacobian\n        \"\"\"\n        value = y\n        log_det_jacobian = 0.0\n        for bijector in reversed(self.bijectors):\n            log_det_jacobian += bijector.batch_log_det_jacobian(value, *args, **kwargs)\n            value = bijector.inverse(value)\n        base_log_pdf = self.base_dist.log_prob(value, *args, **kwargs)\n        if not isinstance(log_det_jacobian, numbers.Number):\n            log_det_jacobian = log_det_jacobian.contiguous().view(*base_log_pdf.size())\n            assert log_det_jacobian.size() == base_log_pdf.size(), \\\n                'Invalid batch_log_det_jacobian().size():\\nexpected {}\\nactual {}'.format(\n                        base_log_pdf.size(), log_det_jacobian.size())\n        return base_log_pdf - log_det_jacobian\n\n\n@add_metaclass(ABCMeta)\nclass Bijector(nn.Module):\n    \"\"\"\n    Abstract class `Bijector`. `Bijector` are bijective transformations with computable\n    log det jacobians. They are meant for use in `TransformedDistribution`.\n    \"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super(Bijector, self).__init__(*args, **kwargs)\n        self.add_inverse_to_cache = False\n\n    @abstractmethod\n    def __call__(self, *args, **kwargs):\n        \"\"\"\n        Virtual forward method\n\n        Invokes the bijection x=>y\n        \"\"\"\n        raise NotImplementedError()\n\n    @abstractmethod\n    def inverse(self, *args, **kwargs):\n        \"\"\"\n        Virtual inverse method\n\n        Inverts the bijection y => x.\n        \"\"\"\n        raise NotImplementedError()\n\n    def log_det_jacobian(self, *args, **kwargs):\n        \"\"\"\n        Default logdet jacobian method.\n\n        Computes the log det jacobian `|dy/dx|`\n        \"\"\"\n        return self.batch_log_det_jacobian(*args, **kwargs).sum()\n\n    def batch_log_det_jacobian(self, *args, **kwargs):\n        \"\"\"\n        Virtual elementwise logdet jacobian method.\n\n        Computes the log abs det jacobian `|dy/dx|`\n        \"\"\"\n        raise NotImplementedError()\n\n\nclass InverseAutoregressiveFlow(Bijector):\n    \"\"\"\n    An implementation of an Inverse Autoregressive Flow. Together with the `TransformedDistribution` this\n    provides a way to create richer variational approximations.\n\n    Example usage::\n\n    >>> base_dist = Normal(...)\n    >>> iaf = InverseAutoregressiveFlow(...)\n    >>> pyro.module(\"my_iaf\", iaf)\n    >>> iaf_dist = TransformedDistribution(base_dist, iaf)\n\n    Note that this implementation is only meant to be used in settings where the inverse of the Bijector\n    is never explicitly computed (rather the result is cached from the forward call). In the context of\n    variational inference, this means that the InverseAutoregressiveFlow should only be used in the guide,\n    i.e. in the variational distribution. In other contexts the inverse could in principle be computed but\n    this would be a (potentially) costly computation that scales with the dimension of the input (and in\n    any case support for this is not included in this implementation).\n\n    :param input_dim: dimension of input\n    :type input_dim: int\n    :param hidden_dim: hidden dimension (number of hidden units)\n    :type hidden_dim: int\n    :param sigmoid_bias: bias on the hidden units fed into the sigmoid; default=`2.0`\n    :type sigmoid_bias: float\n    :param permutation: whether the order of the inputs should be permuted (by default the conditional\n        dependence structure of the autoregression follows the sequential order)\n    :type permutation: bool\n\n    References:\n\n    1. Improving Variational Inference with Inverse Autoregressive Flow [arXiv:1606.04934]\n    Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling\n\n    2. Variational Inference with Normalizing Flows [arXiv:1505.05770]\n    Danilo Jimenez Rezende, Shakir Mohamed\n\n    3. MADE: Masked Autoencoder for Distribution Estimation [arXiv:1502.03509]\n    Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle\n    \"\"\"\n\n    def __init__(self, input_dim, hidden_dim, sigmoid_bias=2.0, permutation=None):\n        super(InverseAutoregressiveFlow, self).__init__()\n        self.input_dim = input_dim\n        self.hidden_dim = hidden_dim\n        self.arn = AutoRegressiveNN(input_dim, hidden_dim, output_dim_multiplier=2, permutation=permutation)\n        self.sigmoid = nn.Sigmoid()\n        self.sigmoid_bias = Variable(torch.Tensor([sigmoid_bias]))\n        self._intermediates_cache = {}\n        self.add_inverse_to_cache = True\n\n    def get_arn(self):\n        \"\"\"\n        :rtype: pyro.nn.AutoRegressiveNN\n\n        Return the AutoRegressiveNN associated with the InverseAutoregressiveFlow\n        \"\"\"\n        return self.arn\n\n    def __call__(self, x, *args, **kwargs):\n        \"\"\"\n        :param x: the input into the bijection\n        :type x: torch.autograd.Variable\n\n        Invokes the bijection x=>y; in the prototypical context of a TransformedDistribution `x` is a\n        sample from the base distribution (or the output of a previous flow)\n        \"\"\"\n        hidden = self.arn(x)\n        sigma = self.sigmoid(hidden[:, 0:self.input_dim] + self.sigmoid_bias.type_as(hidden))\n        mean = hidden[:, self.input_dim:]\n        y = sigma * x + (Variable(torch.ones(sigma.size())).type_as(sigma) - sigma) * mean\n        self._add_intermediate_to_cache(sigma, y, 'sigma')\n        return y\n\n    def inverse(self, y, *args, **kwargs):\n        \"\"\"\n        :param y: the output of the bijection\n        :type y: torch.autograd.Variable\n\n        Inverts y => x. As noted above, this implementation is incapable of inverting arbitrary values\n        `y`; rather it assumes `y` is the result of a previously computed application of the bijector\n        to some `x` (which was cached on the forward call)\n        \"\"\"\n        if (y, 'x') in self._intermediates_cache:\n            x = self._intermediates_cache.pop((y, 'x'))\n            return x\n        else:\n            raise KeyError(\"Bijector InverseAutoregressiveFlow expected to find\" +\n                           \"key in intermediates cache but didn't\")\n\n    def _add_intermediate_to_cache(self, intermediate, y, name):\n        \"\"\"\n        Internal function used to cache intermediate results computed during the forward call\n        \"\"\"\n        assert((y, name) not in self._intermediates_cache),\\\n            \"key collision in _add_intermediate_to_cache\"\n        self._intermediates_cache[(y, name)] = intermediate\n\n    def batch_log_det_jacobian(self, y, *args, **kwargs):\n        \"\"\"\n        Calculates the elementwise determinant of the log jacobian\n        \"\"\"\n        if (y, 'sigma') in self._intermediates_cache:\n            sigma = self._intermediates_cache.pop((y, 'sigma'))\n        else:\n            raise KeyError(\"Bijector InverseAutoregressiveFlow expected to find\" +\n                           \"key in intermediates cache but didn't\")\n        log_sigma = torch.log(sigma)\n        if 'log_pdf_mask' in kwargs:\n            log_sigma = log_sigma * kwargs['log_pdf_mask']\n        return log_sigma\n", "path": "pyro/distributions/transformed_distribution.py"}]}