def test_sample(self):
dim = (1, 128, 128)
parameters = torch.zeros((1, ) + dim)
mfd = nn.MeanFieldDistribution(parameters, dim)
s = mfd.sample()
assert s.shape == parameters.shape
def forward(self, x):
gmrf = nn.MeanFieldDistribution(x, True)
return gmrf.sample()
def test_kl_improper_gmrf(self):
dim = (1, 128, 128)
parameters = torch.zeros((1, ) + dim)
mfd = nn.MeanFieldDistribution(parameters, dim)
kl = mfd.kl_improper_gmrf()
assert np.allclose(kl, np.array([0]))