fixed linting error, changed all map_data tests to use pytest
@@ -7,7 +7,6 @@ import pyro
import pyro.distributions as dist
import pyro.poutine as poutine
from pyro.infer.kl_qp import KL_QP
-from tests.common import TestCase
def eq(x, y, prec=1e-10):
@@ -123,32 +122,29 @@ def test_batch_dim(batch_dim):
assert tr["xs"]["value"].size(1) == data.size(2)
-class NestedMapDataTest(TestCase):
-
- def setUp(self):
-
- self.means = [Variable(torch.randn(2)) for i in range(8)]
- self.mean_batch_size = 2
- self.stds = [Variable(torch.abs(torch.randn(2))) for i in range(6)]
- self.std_batch_size = 3
+def test_nested_map_data():
+ means = [Variable(torch.randn(2)) for i in range(8)]
+ mean_batch_size = 2
+ stds = [Variable(torch.abs(torch.randn(2))) for i in range(6)]
+ std_batch_size = 3
def model(means, stds):
return pyro.map_data("a", means,
- lambda i, x: pyro.map_data("a_{}".format(i), stds,
- lambda j, y: pyro.sample("x_{}{}".format(i, j),
+ lambda i, x:
+ pyro.map_data("a_{}".format(i), stds,
+ lambda j, y:
+ pyro.sample("x_{}{}".format(i, j),
dist.diagnormal, x, y),
- batch_size=self.std_batch_size),
- batch_size=self.mean_batch_size)
+ batch_size=std_batch_size),
+ batch_size=mean_batch_size)
- self.model = model
+ model = model
- def test_default(self):
- xs = self.model(self.means, self.stds)
- self.assertTrue(len(xs) == self.mean_batch_size)
- self.assertTrue(len(xs[0]) == self.std_batch_size)
+ xs = model(means, stds)
+ assert len(xs) == mean_batch_size
+ assert len(xs[0]) == std_batch_size
- def test_trace(self):
- tr = poutine.trace(self.model)(self.means, self.stds)
+ tr = poutine.trace(model)(means, stds)
for name in tr.keys():
if tr[name]["type"] == "sample":
- self.assertTrue(tr[name]["scale"] == 4.0 * 2.0)
+ assert tr[name]["scale"] == 4.0 * 2.0
