Disagreement in log probabilities between log_prob and sample_and_log_prob

[cbames]

I'm seeing an odd behavior when contexts are included. If I compare the log probabilities output by sample_and_log_prob with the log probabilities calculated by log_prob on the samples returned by sample_and_log_prob I get different results:

log_prob: [[-16.1957, -16.7197, -17.4852, -20.2420, -17.6908]]
log_prob_2: [-16.1218, -16.7461, -16.9846, -20.4161, -18.0095]

I think I've tracked it down somewhat to the fact that the noise produced in sample_and_log_prob does not match the reconstruction of the noise in log_prob.

Noise sampled in sample_and_log_prob:

noise in s and l: tensor([[[-1.4298,  2.0084, -0.6241,  0.3967,  0.5529,  0.4732, -0.8063],
         [-0.2945, -1.4018, -0.1627, -0.1684, -1.3888,  0.2485, -0.3683],
         [-0.3361, -0.1749,  0.2021,  0.1011,  0.9791, -2.1958, -0.8109],
         [ 0.1006,  2.5615, -0.0782,  1.9179,  1.6321, -0.7352,  0.5438],
         [ 1.4222, -0.1466, -1.3136,  0.9655, -0.3346,  0.8428,  0.0655]]])

Noise reconstructed in log_prob:

noise in log: tensor([[-1.5227,  1.9867, -0.6582,  0.3719,  0.5637,  0.3799, -0.5922],
        [-0.3344, -1.4428, -0.1877, -0.1901, -1.3203,  0.3617, -0.2444],
        [-0.3513, -0.2244,  0.1518,  0.0817,  1.0594, -1.9549, -0.6088],
        [ 0.2416,  2.5659, -0.1450,  1.9957,  1.6983, -0.4731,  0.6077],
        [ 1.5897, -0.1721, -1.2935,  1.0161, -0.2590,  0.8674,  0.1447]],
       grad_fn=<AddBackward0>)

If I run the same test without contexts there is no discrepancy between the log probabilities. I've included my test case as well:

import torch 
from nflows.flows.base import Flow
from nflows.distributions.normal import StandardNormal
from nflows.transforms.base import CompositeTransform
from nflows.transforms.autoregressive import MaskedAffineAutoregressiveTransform
from nflows.transforms.permutations import RandomPermutation 
from nflows.transforms.lu import LULinear
from nflows.nn.nets import ResidualNet

#Model definition 
base_dist = StandardNormal(shape=[7])

maf_transforms = []
for i in range(3):
    maf_transforms.append(RandomPermutation(features=7))
    maf_transforms.append(MaskedAffineAutoregressiveTransform(features=7,
                                                              hidden_features = 256, 
                                                              context_features=8))

maf_transform = CompositeTransform(maf_transforms)
maf_flow = Flow(maf_transform, base_dist)

test = torch.FloatTensor([0]*8)

samples, log_prob = maf_flow.sample_and_log_prob(5, context=test.unsqueeze(0))

log_prob_2 = maf_flow.log_prob(samples.squeeze())

num_close = torch.sum(torch.abs(log_prob - log_prob_2) < 1e-2)

print("log_prob:", log_prob)
print("log_prob_2:", log_prob_2)

print("num close:", num_close)

EDIT

nvm, i see I missed including the context with the log_prob. Tough to calculate the proper log probability without the context.

[cbames]

api misuse.