import numpy as np
import torch
from torch.autograd import Variable

from pyro.distributions.distribution import Distribution


class LogNormal(Distribution):
    """
    :param mu: mean *(vector)*
    :param sigma: standard deviations *(vector (0, Infinity))*

    A distribution over probability vectors obtained by exp-transforming a random
    variable drawn from ``Normal({mu: mu, sigma: sigma})``.
    """
    reparameterized = True

    def __init__(self, mu, sigma, batch_size=None, *args, **kwargs):
        """
        Params:
          `mu` - mean
          `sigma` - root variance
        """
        self.mu = mu
        self.sigma = sigma
        if mu.size() != sigma.size():
            raise ValueError("Expected mu.size() == sigma.size(), but got {} vs {}"
                             .format(mu.size(), sigma.size()))
        if mu.dim() == 1 and batch_size is not None:
            self.mu = mu.expand(batch_size, mu.size(0))
            self.sigma = sigma.expand(batch_size, sigma.size(0))
        super(LogNormal, self).__init__(*args, **kwargs)

    def batch_shape(self, x=None):
        event_dim = 1
        mu = self.mu
        if x is not None and x.size() != mu.size():
            mu = self.mu.expand(x.size()[:-event_dim] + self.event_shape())
        return mu.size()[:-event_dim]

    def event_shape(self):
        event_dim = 1
        return self.mu.size()[-event_dim:]

    def shape(self, x=None):
        return self.batch_shape(x) + self.event_shape()

    def sample(self):
        """
        Reparameterized log-normal sampler.
        """
        eps = Variable(torch.randn(1).type_as(self.mu.data))
        z = self.mu + self.sigma * eps
        return torch.exp(z)

    def batch_log_pdf(self, x):
        """
        log-normal log-likelihood
        """
        mu = self.mu.expand(self.shape(x))
        sigma = self.sigma.expand(self.shape(x))
        ll_1 = Variable(torch.Tensor([-0.5 * np.log(2.0 * np.pi)])
                        .type_as(mu.data).expand_as(x))
        ll_2 = -torch.log(sigma * x)
        ll_3 = -0.5 * torch.pow((torch.log(x) - mu) / sigma, 2.0)
        batch_log_pdf = torch.sum(ll_1 + ll_2 + ll_3, -1)
        batch_log_pdf_shape = self.batch_shape(x) + (1,)
        return batch_log_pdf.contiguous().view(batch_log_pdf_shape)

    def analytic_mean(self):
        return torch.exp(self.mu + 0.5 * torch.pow(self.sigma, 2.0))

    def analytic_var(self):
        return (torch.exp(torch.pow(self.sigma, 2.0)) - Variable(torch.ones(1))) * \
            torch.pow(self.analytic_mean(), 2)

pyro/distributions/log_normal.py
import numpy as np
import torch
from torch.autograd import Variable

from pyro.distributions.distribution import Distribution


class LogNormal(Distribution):
    """
    :param mu: mean *(vector)*
    :param sigma: standard deviations *(vector (0, Infinity))*

    A distribution over probability vectors obtained by exp-transforming a random
    variable drawn from ``Normal({mu: mu, sigma: sigma})``.
    """
    reparameterized = True

    def __init__(self, mu, sigma, batch_size=None, *args, **kwargs):
        """
        Params:
          `mu` - mean
          `sigma` - root variance
        """
        self.mu = mu
        self.sigma = sigma
        if mu.size() != sigma.size():
            raise ValueError("Expected mu.size() == sigma.size(), but got {} vs {}"
                             .format(mu.size(), sigma.size()))
        if mu.dim() == 1 and batch_size is not None:
            self.mu = mu.expand(batch_size, mu.size(0))
            self.sigma = sigma.expand(batch_size, sigma.size(0))
        super(LogNormal, self).__init__(*args, **kwargs)

    def batch_shape(self, x=None):
        event_dim = 1
        mu = self.mu
        if x is not None and x.size() != mu.size():
            mu = self.mu.expand(x.size()[:-event_dim] + self.event_shape())
        return mu.size()[:-event_dim]

    def event_shape(self):
        event_dim = 1
        return self.mu.size()[-event_dim:]

    def shape(self, x=None):
        return self.batch_shape(x) + self.event_shape()

    def sample(self):
        """
        Reparameterized log-normal sampler.
        """
        eps = Variable(torch.randn(self.mu.size()).type_as(self.mu.data))
        z = self.mu + self.sigma * eps
        return torch.exp(z)

    def batch_log_pdf(self, x):
        """
        log-normal log-likelihood
        """
        mu = self.mu.expand(self.shape(x))
        sigma = self.sigma.expand(self.shape(x))
        ll_1 = Variable(torch.Tensor([-0.5 * np.log(2.0 * np.pi)])
                        .type_as(mu.data).expand_as(x))
        ll_2 = -torch.log(sigma * x)
        ll_3 = -0.5 * torch.pow((torch.log(x) - mu) / sigma, 2.0)
        batch_log_pdf = torch.sum(ll_1 + ll_2 + ll_3, -1)
        batch_log_pdf_shape = self.batch_shape(x) + (1,)
        return batch_log_pdf.contiguous().view(batch_log_pdf_shape)

    def analytic_mean(self):
        return torch.exp(self.mu + 0.5 * torch.pow(self.sigma, 2.0))

    def analytic_var(self):
        return (torch.exp(torch.pow(self.sigma, 2.0)) - Variable(torch.ones(1))) * \
            torch.pow(self.analytic_mean(), 2)

pyro/distributions/log_normal.py
diff --git a/pyro/distributions/log_normal.py b/pyro/distributions/log_normal.py
index 76d65c7342..2c8dc23c63 100644
--- a/pyro/distributions/log_normal.py
+++ b/pyro/distributions/log_normal.py
@@ -49,7 +49,7 @@ def sample(self):
         """
         Reparameterized log-normal sampler.
         """
-        eps = Variable(torch.randn(1).type_as(self.mu.data))
+        eps = Variable(torch.randn(self.mu.size()).type_as(self.mu.data))
         z = self.mu + self.sigma * eps
         return torch.exp(z)
 
diff --git a/tests/distributions/dist_fixture.py b/tests/distributions/dist_fixture.py
index e48fab4b5b..bdb895fc77 100644
--- a/tests/distributions/dist_fixture.py
+++ b/tests/distributions/dist_fixture.py
@@ -39,7 +39,7 @@ def get_num_test_data(self):
         return len(self.test_data)
 
     def get_samples(self, num_samples, **dist_params):
-        return [self.pyro_dist(**dist_params).data.cpu().numpy() for _ in range(num_samples)]
+        return self.pyro_dist(batch_size=num_samples, **dist_params)
 
     def get_test_data(self, idx, wrap_tensor=True):
         if not wrap_tensor:
diff --git a/tests/distributions/test_distributions.py b/tests/distributions/test_distributions.py
index bbbb310efc..1bf3d5c253 100644
--- a/tests/distributions/test_distributions.py
+++ b/tests/distributions/test_distributions.py
@@ -94,11 +94,11 @@ def test_mean_and_variance(dist):
     num_samples = dist.get_num_samples(SINGLE_TEST_DATUM_IDX)
     dist_params = dist.get_dist_params(SINGLE_TEST_DATUM_IDX)
     torch_samples = dist.get_samples(num_samples, **dist_params)
-    sample_mean = np.mean(torch_samples, 0)
-    sample_var = np.var(torch_samples, 0)
+    sample_mean = torch_samples.float().mean(0)
+    sample_var = torch_samples.float().var(0)
     try:
-        analytic_mean = unwrap_variable(dist.pyro_dist.analytic_mean(**dist_params))
-        analytic_var = unwrap_variable(dist.pyro_dist.analytic_var(**dist_params))
+        analytic_mean = dist.pyro_dist.analytic_mean(**dist_params)
+        analytic_var = dist.pyro_dist.analytic_var(**dist_params)
         assert_equal(sample_mean, analytic_mean, prec=dist.prec)
         assert_equal(sample_var, analytic_var, prec=dist.prec)
     except (NotImplementedError, ValueError):

Test batch_size parameter is correctly applied for all distributions
With #323, we are now supporting `batch_size` parameter for both the class and functional forms of both the distributions. Let us add a simple test to validate that this is working as expected. 
